{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Notebook to demonstrate TAO workflow on purpose built models\n",
    "\n",
    "Transfer learning is the process of transferring learned features from one application to another. It is a commonly used training technique where you use a model trained on one task and re-train to use it on a different task. Train Adapt Optimize (TAO) Toolkit  is a simple and easy-to-use Python based AI toolkit for taking purpose-built AI models and customizing them with users' own data.\n",
    "\n",
    "![image](https://d29g4g2dyqv443.cloudfront.net/sites/default/files/akamai/TAO/tlt-tao-toolkit-bring-your-own-model-diagram.png)\n",
    "\n",
    "### The workflow in a nutshell\n",
    "\n",
    "- Creating a dataset\n",
    "- Upload dataset to the service\n",
    "- Running dataset convert (for specific models)\n",
    "- Getting a PTM from NGC\n",
    "- Model Actions\n",
    "    - Train (Normal/AutoML)\n",
    "    - Evaluate\n",
    "    - Prune, retrain (for specific models)\n",
    "    - Export\n",
    "    - TAO-Deploy (for specific models)\n",
    "    - Inference on TAO\n",
    "    - Inference on TRT (for specific models)\n",
    "    \n",
    "### Table of contents\n",
    "\n",
    "1. [Create datasets](#head-1)\n",
    "1. [List the created datasets](#head-2)\n",
    "1. [Dataset convert Action for train dataset](#head-3) (for specific models)\n",
    "1. [Dataset convert Action for val dataset](#head-3.1) (for specific models)\n",
    "1. [Create an experiment](#head-4)\n",
    "1. [List experiments](#head-5)\n",
    "1. [Assign datasets](#head-6)\n",
    "1. [Assign PTM](#head-7)\n",
    "1. [View hyperparameters that are enabled by default](#head-8)\n",
    "1. [Set AutoML related configurations](#head-9)\n",
    "1. [Actions](#head-10)\n",
    "1. [Train](#head-11)\n",
    "1. [Evaluate](#head-12)\n",
    "1. [Optimize: Apply specs for prune](#head-14) (for specific models)   \n",
    "1. [Optimize: Apply specs for retrain](#head-15) (for specific models)\n",
    "1. [Optimize: Run actions](#head-16) (for specific models)\n",
    "1. [Export](#head-17)\n",
    "1. [TRT Engine generation using TAO-Deploy](#head-18) (for specific models)\n",
    "1. [TAO inference](#head-19)\n",
    "1. [TRT inference](#head-20) (for specific models)\n",
    "\n",
    "### Requirements\n",
    "Please find the server requirements [here](https://docs.nvidia.com/tao/tao-toolkit/text/tao_toolkit_api/api_setup.html#)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import json\n",
    "import os\n",
    "import requests\n",
    "import shutil\n",
    "import time\n",
    "from IPython.display import clear_output\n",
    "import subprocess\n",
    "import glob"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### FIXME\n",
    "\n",
    "1. Assign a model_name in FIXME 1\n",
    "\n",
    "    1.1 Assign model type for action_recognition/fpenet/lprnet/pose_classification in FIXME 1.1\n",
    "\n",
    "    1.2 Assign platform for action_recognition in FIXME 1.2\n",
    "    \n",
    "    1.3 Assign model input type for action_recognition in FIXME 1.3\n",
    "1. Assign a workdir in FIXME 2\n",
    "1. Assign the ip_address and port_number in FIXME 3 ([info](https://docs.nvidia.com/tao/tao-toolkit/text/tao_toolkit_api/api_rest_api.html))\n",
    "1. Assign the ngc_api_key variable in FIXME 4\n",
    "1. (Optional) Enable AutoML if needed in FIXME 5\n",
    "1. (Optional) Choose between bayesian and hyperband automl_algorithm in FIXME 6 (If automl was enabled in FIXME5)\n",
    "1. Choose to download jobs or not in FIXME 7\n",
    "1. Choose between default and custom dataset in FIXME 8\n",
    "1. Assign path of DATA_DIR in FIXME 9"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Define model_name workspaces and other variables\n",
    "# Available models (#FIXME 1):\n",
    "# 1. action_recognition - https://docs.nvidia.com/tao/tao-toolkit/text/action_recognition_net.html\n",
    "# 2. bpnet - https://docs.nvidia.com/tao/tao-toolkit/text/bodypose_estimation/bodyposenet.html\n",
    "# 3. fpenet - https://docs.nvidia.com/tao/tao-toolkit/text/facial_landmarks_estimation/facial_landmarks_estimation.html\n",
    "# 4. lprnet - https://docs.nvidia.com/tao/tao-toolkit/text/character_recognition/index.html\n",
    "# 5. ml_recog - https://docs.nvidia.com/tao/tao-toolkit/text/ml_recog/index.html\n",
    "# 6. ocdnet - https://docs.nvidia.com/tao/tao-toolkit/text/ocdnet/index.html\n",
    "# 7. ocrnet - https://docs.nvidia.com/tao/tao-toolkit/text/ocrnet/index.html\n",
    "# 8. optical_inspection - https://docs.nvidia.com/tao/tao-toolkit/text/optical_inspection/index.html\n",
    "# 9. pose_classification - https://docs.nvidia.com/tao/tao-toolkit/text/pose_classification/index.html\n",
    "# 10. pointpillars - https://docs.nvidia.com/tao/tao-toolkit/text/point_cloud/pointpillars.html\n",
    "# 11. re_identification - https://docs.nvidia.com/tao/tao-toolkit/text/re_identification/index.html\n",
    "# 12. centerpose -\n",
    "# 13. visual_changenet_segment - https://docs.nvidia.com/tao/tao-toolkit/text/visual_changenet/index.html\n",
    "# 14. visual_changenet_classify - https://docs.nvidia.com/tao/tao-toolkit/text/visual_changenet/index.html \n",
    "\n",
    "model_name = \"action_recognition\" # FIXME1 (Add the model name from the above mentioned list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if model_name in (\"action_recognition\",\"fpenet\",\"lprnet\",\"pose_classification\"):\n",
    "    # FIXME1.1 - model_type - string\n",
    "        # action-recognition: rgb/of/joint;\n",
    "        # fpenet: 10/80 (value represents the number of keypoints)\n",
    "        # lprnet: us/ch (us for United States, ch for China)\n",
    "        # pose-classification: kinetics/nvidia\n",
    "    model_type = \"rgb\"\n",
    "\n",
    "    if model_name == \"action_recognition\":\n",
    "        if model_type not in (\"rgb\",\"of\",\"joint\"):\n",
    "            raise Exception(\"Choose one of rgb/of/joint for action recognition model_type\")\n",
    "    elif model_name == \"fpenet\":\n",
    "        if model_type not in (\"10\",\"80\"):\n",
    "            raise Exception(\"Choose one of 10/80 for FPENET model_type\")\n",
    "    elif model_name == \"lprnet\":\n",
    "        if model_type not in (\"us\",\"ch\"):\n",
    "            raise Exception(\"Choose one of us/ch for LPRNET model_type\")\n",
    "    elif model_name == \"pose_classification\":\n",
    "        if model_type not in (\"kinetics\",\"nvidia\"):\n",
    "            raise Exception(\"Choose one of kinetics/nvidia for pose classification model_type\")\n",
    "\n",
    "    if model_name == \"action_recognition\":\n",
    "        platform = \"a100\" # FIXME1.2 a100/xavier - valid only for model_type that is not rgb\n",
    "        model_input_type = \"3d\" # FIXME1.3 3d/2d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "workdir = \"workdir_purpose_built_models\" # FIXME2\n",
    "host_url = \"http://<ip_address>:<port_number>\" # FIXME3 example: https://10.137.149.22:32334\n",
    "# In host machine, node ip_address and port number can be obtained as follows,\n",
    "# ip_address: hostname -i\n",
    "# port_number: kubectl get service ingress-nginx-controller -o jsonpath='{.spec.ports[0].nodePort}'\n",
    "ngc_api_key = \"<ngc_api_key>\" # FIXME4 example: (Add NGC API key)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "automl_enabled = False # FIXME5 set to True if you want to run automl for the model chosen in the previous cell\n",
    "automl_algorithm = \"bayesian\" # FIXME6 example: bayesian/hyperband\n",
    "# FIXME7 Defaulted to False as downloading jobs from service to your machine takes time\n",
    "# Set to True if you want to download jobs where examples have been provided like for train, export, inference.\n",
    "download_jobs = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Exchange NGC_API_KEY for JWT\n",
    "data = json.dumps({\"ngc_api_key\": ngc_api_key})\n",
    "response = requests.post(f\"{host_url}/api/v1/login\", data=data)\n",
    "assert response.status_code in (200, 201)\n",
    "assert \"user_id\" in response.json().keys()\n",
    "user_id = response.json()[\"user_id\"]\n",
    "print(\"User ID\",user_id)\n",
    "assert \"token\" in response.json().keys()\n",
    "token = response.json()[\"token\"]\n",
    "print(\"JWT\",token)\n",
    "\n",
    "# Set base URL\n",
    "base_url = f\"{host_url}/api/v1/users/{user_id}\"\n",
    "print(\"API Calls will be forwarded to\",base_url)\n",
    "\n",
    "headers = {\"Authorization\": f\"Bearer {token}\"}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Creating workdir\n",
    "if not os.path.isdir(workdir):\n",
    "    os.makedirs(workdir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Function to split tar files <a class=\"anchor\" id=\"head-1.1\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import tarfile\n",
    "\n",
    "def split_tar_file(input_tar_path, output_dir, max_split_size=0.2*1024*1024*1024):\n",
    "\tos.makedirs(output_dir, exist_ok=True)\n",
    "\t\n",
    "\twith tarfile.open(input_tar_path, 'r') as original_tar:\n",
    "\t\tmembers = original_tar.getmembers()\n",
    "\t\tcurrent_split_size = 0\n",
    "\t\tcurrent_split_number = 0\n",
    "\t\tcurrent_split_name = os.path.join(output_dir, f'smaller_file_{current_split_number}.tar')\n",
    "\t\t\n",
    "\t\twith tarfile.open(current_split_name, 'w') as split_tar:\n",
    "\t\t\tfor member in members:\n",
    "\t\t\t\tif current_split_size + member.size <= max_split_size:\n",
    "\t\t\t\t\tsplit_tar.addfile(member, original_tar.extractfile(member))\n",
    "\t\t\t\t\tcurrent_split_size += member.size\n",
    "\t\t\t\telse:\n",
    "\t\t\t\t\tsplit_tar.close()\n",
    "\t\t\t\t\tcurrent_split_number += 1\n",
    "\t\t\t\t\tcurrent_split_name = os.path.join(output_dir, f'smaller_file_{current_split_number}.tar')\n",
    "\t\t\t\t\tcurrent_split_size = 0\n",
    "\t\t\t\t\tsplit_tar = tarfile.open(current_split_name, 'w')  # Open a new split tar archive\n",
    "\t\t\t\t\tsplit_tar.addfile(member, original_tar.extractfile(member))\n",
    "\t\t\t\t\tcurrent_split_size += member.size"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set dataset type, format <a class=\"anchor\" id=\"head-1.1\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Action Recognition:** We will be using the HMDB51 [dataset](https://serre-lab.clps.brown.edu/resource/hmdb-a-large-human-motion-database/) for the tutorial. (We choose catch/smile for this tutorial):\n",
    "\n",
    "**BPNET:** We will be using the `COCO dataset` for Instance segmentation - MaskRCNN. `download_coco.sh` script from dataset prepare will be used to download and unzip the coco2017 dataset from [here](https://cocodataset.org/#download)\n",
    "\n",
    "**FPENET:** We will be using `AFW dataset`. Download it from [here](https://ibug.doc.ic.ac.uk/download/annotations/afw.zip/) and place it in $DATA_DIR.\n",
    "\n",
    "**LPRNET**: We will be using the `OpenALPR benchmark dataset` for the tutorial. The following script will download the dataset automatically and convert it to the format used by TAO.  \n",
    "\n",
    "**MLRecogNet** We will be using the `Retail Product Checkout Dataset` for the tutorial. Downdload the datsaet from [here](https://www.kaggle.com/datasets/diyer22/retail-product-checkout-dataset) and place it under $DATA_DIR/metric_learning_recognition\n",
    "\n",
    "**OCDNET**: We will be using the ICDAR2015 dataset for the ocdnet tutorial. Please access the dataset [here](https://rrc.cvc.uab.es/?ch=4&com=tasks) to register and download the data from Task 4.1: Text Localization. Unzip the files to DATA_DIR\n",
    "\n",
    "**OCRNET**: We will be using the ICDAR15 word recognition dataset for the tutorial. To find more details please visit [here](\n",
    "https://rrc.cvc.uab.es/?ch=4&com=tasks). Please download the ICDAR15 word recognition train dataset and test_dataset [here](https://rrc.cvc.uab.es/?ch=4&com=downloads) to DATA_DIR.\n",
    "\n",
    "**Pointpillars:** We will be using the `kitti object detection dataset` for this example. To find more details, please visit [here](http://www.cvlibs.net/datasets/kitti/eval_object.php?obj_benchmark=2d)\n",
    "\n",
    "**Pose Classification:** We will be using the Kinetics dataset from [Deepmind](https://deepmind.com/research/open-source/kinetics) or NVIDIA created dataset. For kinetics based dataset set model_type as `kinetics` and for nvidia based dataset set model_type as `nvidia`\n",
    "\n",
    "**Re-Identification:** We will be using the [Market-1501](https://zheng-lab.cecs.anu.edu.au/Project/project_reid.html) dataset. Download the dataset [here](https://drive.google.com/file/d/1TwkgQcIa_EgRjVMPSbyEKtcfljqURrzi/view?usp=sharing) and extract it.\n",
    "\n",
    "**Optical Inspection:** Bring your own dataset according to the format described [here](https://docs.nvidia.com/tao/tao-toolkit/text/data_annotation_format.html#optical-inspection-format). \n",
    "\n",
    "**Visual ChangeNet-Classification:** Bring your own dataset according to the format described [here](https://docs.nvidia.com/tao/tao-toolkit/text/data_annotation_format.html#optical-inspection-format). \n",
    "\n",
    "**Visual ChangeNet-Segmentation:** We will be using the [Market-1501](https://zheng-lab.cecs.anu.edu.au/Project/project_reid.html) dataset. Download the dataset [here](https://www.dropbox.com/s/18fb5jo0npu5evm/LEVIR-CD256.zip) and extract it. \n",
    "\n",
    "**CenterPose:** We will be using [Google Objectron](https://github.com/google-research-datasets/Objectron) dataset. The following script will download and preprocess the dataset the dataset automatically."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if model_name == \"lprnet\":\n",
    "    ds_type = \"character_recognition\"\n",
    "    ds_format = \"lprnet\"\n",
    "elif model_name in (\"visual_changenet_classify\", \"visual_changenet_segment\"):\n",
    "    ds_format = model_name\n",
    "    ds_type = model_name = \"visual_changenet\"\n",
    "else:\n",
    "    ds_type = model_name\n",
    "    ds_format = \"default\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "dataset_to_be_used = \"default\" #FIXME8 #default/custom; default for the dataset used in this tutorial notebook; custom for a different dataset\n",
    "DATA_DIR = os.path.abspath(model_name) # FIXME9 (set absolute path of the data_directory)\n",
    "os.environ['DATA_DIR']= DATA_DIR\n",
    "!mkdir -p $DATA_DIR\n",
    "job_map = {}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Dataset download and pre-processing <a class=\"anchor\" id=\"head-1\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if dataset_to_be_used == \"default\":\n",
    "    if model_name == \"action_recognition\":\n",
    "        !sudo apt-get update -y && sudo apt-get install unrar-free -y\n",
    "        !wget -P $DATA_DIR --no-check-certificate http://serre-lab.clps.brown.edu/wp-content/uploads/2013/10/hmdb51_org.rar\n",
    "        assert os.path.exists(f\"{DATA_DIR}/hmdb51_org.rar\")\n",
    "        !mkdir -p $DATA_DIR/videos && unrar x -o+ $DATA_DIR/hmdb51_org.rar $DATA_DIR/videos\n",
    "        !mkdir -p $DATA_DIR/raw_data\n",
    "        !unrar x -o+ $DATA_DIR/videos/catch.rar $DATA_DIR/raw_data\n",
    "        !unrar x -o+ $DATA_DIR/videos/smile.rar $DATA_DIR/raw_data\n",
    "        assert os.path.exists(f\"{DATA_DIR}/raw_data/catch\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/raw_data/smile\")\n",
    "    elif model_name == \"bpnet\":\n",
    "        !bash ../dataset_prepare/coco/download_coco.sh $DATA_DIR\n",
    "        # Remove existing data\n",
    "        !rm -rf $DATA_DIR/train2017/images\n",
    "        !rm -rf $DATA_DIR/val2017/images\n",
    "        # Rearrange data in the required format\n",
    "        !mv $DATA_DIR/raw-data/* $DATA_DIR/\n",
    "        !cp ../dataset_prepare/bpnet/* $DATA_DIR/\n",
    "        assert os.path.exists(f\"{DATA_DIR}/train2017\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/val2017\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/annotations\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/bpnet_18joints.json\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/coco_spec.json\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/infer_spec.yaml\")\n",
    "    elif model_name == \"fpenet\":\n",
    "        assert os.path.exists(f\"{DATA_DIR}/afw.zip\")\n",
    "        !mkdir $DATA_DIR/data\n",
    "        !unzip -uq $DATA_DIR/afw.zip -d $DATA_DIR/data/afw\n",
    "        !cp ../dataset_prepare/fpenet/data.json $DATA_DIR/\n",
    "        assert os.path.exists(f\"{DATA_DIR}/data/afw\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/data.json\")\n",
    "    elif model_name == \"lprnet\":\n",
    "        !python3 -m pip install --upgrade pip\n",
    "        !python3 -m pip install \"opencv-python>=3.4.0.12,<=4.5.5.64\"\n",
    "        !bash ../dataset_prepare/lprnet/download_and_prepare_data.sh $DATA_DIR\n",
    "        assert os.path.exists(f\"{DATA_DIR}/train/image\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/train/label\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/val/image\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/val/label\")\n",
    "    elif model_name == \"ml_recog\":\n",
    "        assert os.path.exists(f\"{DATA_DIR}/metric_learning_recognition/retail-product-checkout-dataset.zip\")\n",
    "        !unzip -uq $DATA_DIR/metric_learning_recognition/retail-product-checkout-dataset.zip -d $DATA_DIR/metric_learning_recognition\n",
    "    elif model_name == \"ocdnet\":\n",
    "        assert(os.path.exists(f\"{DATA_DIR}/train/img\"))\n",
    "        assert(os.path.exists(f\"{DATA_DIR}/train/gt\"))\n",
    "        assert(os.path.exists(f\"{DATA_DIR}/test/img\"))\n",
    "        assert(os.path.exists(f\"{DATA_DIR}/test/gt\"))\n",
    "    elif model_name == \"ocrnet\":\n",
    "        !mkdir -p $DATA_DIR/train && rm -rf $DATA_DIR/train/*\n",
    "        !mkdir -p $DATA_DIR/test && rm -rf $DATA_DIR/test/*\n",
    "        !unzip -u $DATA_DIR/ch4_test_word_images_gt.zip -d $DATA_DIR/test\n",
    "        !cp $DATA_DIR/Challenge4_Test_Task3_GT.txt -d $DATA_DIR/test\n",
    "        !unzip -u $DATA_DIR/ch4_training_word_images_gt.zip -d $DATA_DIR/train    \n",
    "        assert os.path.exists(f\"{DATA_DIR}/ch4_test_word_images_gt.zip\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/Challenge4_Test_Task3_GT.txt\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/ch4_training_word_images_gt.zip\")\n",
    "    elif model_name == \"optical_inspection\" or ds_format == \"visual_changenet_classify\":\n",
    "        assert os.path.exists(f\"{DATA_DIR}/train/images\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/train/dataset.csv\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/val/images\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/val/dataset.csv\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/test/images\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/test/dataset.csv\")\n",
    "    elif model_name == \"visual_changenet\" and ds_format == \"visual_changenet_segment\":\n",
    "        #Download the data\n",
    "        URL_DATASET = \"https://www.dropbox.com/s/18fb5jo0npu5evm/LEVIR-CD256.zip\"\n",
    "        os.environ[\"URL_DATASET\"]=URL_DATASET\n",
    "        !if [ ! -f $DATA_DIR/LEVIR-CD256.zip ]; then wget $URL_DATASET -O $DATA_DIR/LEVIR-CD-256.zip; else echo \"image archive already downloaded\"; fi \n",
    "        # Check the dataset is present\n",
    "        !mkdir -p $DATA_DIR\n",
    "        !if [ ! -f $DATA_DIR/LEVIR-CD-256.zip ]; then echo 'Dataset zip file not found, please download.'; else echo 'Found Dataset zip file.';fi\n",
    "        # unpack \n",
    "        !unzip -u $DATA_DIR/LEVIR-CD-256.zip -d $DATA_DIR\n",
    "        assert os.path.exists(f\"{DATA_DIR}/LEVIR-CD256/A\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/LEVIR-CD256/B\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/LEVIR-CD256/label\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/LEVIR-CD256/list/train.txt\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/LEVIR-CD256/list/val.txt\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/LEVIR-CD256/list/test.txt\")\n",
    "    elif model_name == \"pointpillars\":\n",
    "        !unzip -u $DATA_DIR/data_object_image_2.zip -d $DATA_DIR\n",
    "        !unzip -u $DATA_DIR/data_object_label_2.zip -d $DATA_DIR\n",
    "        !unzip -u $DATA_DIR/data_object_velodyne.zip -d $DATA_DIR\n",
    "        !unzip -u $DATA_DIR/data_object_calib.zip -d $DATA_DIR\n",
    "        assert os.path.exists(f\"{DATA_DIR}/training/image_2\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/training/label_2\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/training/velodyne\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/training/calib\")\n",
    "    elif model_name == \"pose_classification\":\n",
    "        !pip3 install -U gdown\n",
    "        if model_type == \"kinetics\":\n",
    "            !gdown https://drive.google.com/uc?id=1dmzCRQsFXJ18BlXj1G9sbDnsclXIdDdR -O $DATA_DIR/st-gcn-processed-data.zip\n",
    "            !unzip $DATA_DIR/st-gcn-processed-data.zip -d $DATA_DIR\n",
    "            !mv $DATA_DIR/data/Kinetics/kinetics-skeleton $DATA_DIR/kinetics\n",
    "            !rm -r $DATA_DIR/data\n",
    "            !rm $DATA_DIR/st-gcn-processed-data.zip\n",
    "            assert os.path.exists(f\"{DATA_DIR}/kinetics\")\n",
    "        elif model_type == \"nvidia\":\n",
    "            !gdown https://drive.google.com/uc?id=1GhSt53-7MlFfauEZ2YkuzOaZVNIGo_c- -O $DATA_DIR/data_3dbp_nvidia.zip\n",
    "            !mkdir -p $DATA_DIR/nvidia\n",
    "            !unzip $DATA_DIR/data_3dbp_nvidia.zip -d $DATA_DIR/nvidia\n",
    "            !rm $DATA_DIR/data_3dbp_nvidia.zip\n",
    "            assert os.path.exists(f\"{DATA_DIR}/nvidia\")\n",
    "            assert os.path.exists(f\"{DATA_DIR}/{model_type}/train_data.npy\") and os.path.exists(f\"{DATA_DIR}/{model_type}/train_label.pkl\") and os.path.exists(f\"{DATA_DIR}/{model_type}/val_data.npy\") and os.path.exists(f\"{DATA_DIR}/{model_type}/val_label.pkl\")\n",
    "    elif model_name == \"re_identification\":\n",
    "        !pip3 install -U gdown\n",
    "        !gdown https://drive.google.com/uc?id=0B8-rUzbwVRk0c054eEozWG9COHM -O $DATA_DIR/market1501.zip\n",
    "        !unzip -u $DATA_DIR/market1501.zip -d $DATA_DIR\n",
    "        !rm -rf $DATA_DIR/market1501\n",
    "        !mv $DATA_DIR/Market-1501-v15.09.15 $DATA_DIR/market1501\n",
    "        !rm $DATA_DIR/market1501.zip\n",
    "        assert os.path.exists(f\"{DATA_DIR}/market1501\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if model_name in (\"lprnet\",\"ocdnet\",\"ocrnet\", \"optical_inspection\") or ds_format in (\"visual_changenet_classify\"):\n",
    "    eval_dataset_path = f\"{DATA_DIR}/purpose_built_models_val.tar.gz\"\n",
    "if model_name in (\"lprnet\", \"optical_inspection\") or ds_format in (\"visual_changenet_classify\"):\n",
    "    test_dataset_path = f\"{DATA_DIR}/purpose_built_models_test.tar.gz\"\n",
    "train_dataset_path = f\"{DATA_DIR}/purpose_built_models_train.tar.gz\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Create train dataset\n",
    "data = json.dumps({\"type\":ds_type,\"format\":ds_format})\n",
    "endpoint = f\"{base_url}/datasets\"\n",
    "response = requests.post(endpoint,data=data,headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "\n",
    "print(response)\n",
    "print(response.json())\n",
    "assert \"id\" in response.json().keys()\n",
    "train_dataset_id = response.json()[\"id\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "if dataset_to_be_used == \"default\":\n",
    "    USER_EXPERIMENT_DIR = os.path.join(\"/shared/users\",user_id,\"datasets\",train_dataset_id)\n",
    "    if model_name == \"action_recognition\":\n",
    "        !python3 -m pip install opencv-python numpy\n",
    "        # For rgb action recognition\n",
    "        !if [ -d tao_toolkit_recipes ]; then rm -rf tao_toolkit_recipes; fi\n",
    "        !git clone https://github.com/NVIDIA-AI-IOT/tao_toolkit_recipes\n",
    "        assert os.path.exists(\"tao_toolkit_recipes\")\n",
    "        !cd tao_toolkit_recipes/tao_action_recognition/data_generation/ && bash ./preprocess_HMDB_RGB.sh $DATA_DIR/raw_data $DATA_DIR/processed_data \n",
    "\n",
    "        # For optical flow, comment the above 3 lines and uncomment the below (Note: for generating optical flow, a Turing or Ampere above GPU is needed.)\n",
    "        #!echo <passwd> | sudo -S apt install -y libfreeimage-dev\n",
    "        #!cp ../dataset_prepare/action_recognition/AppOFCuda tao_toolkit_recipes/tao_action_recognition/data_generation/\n",
    "        #!cd tao_toolkit_recipes/tao_action_recognition/data_generation/ && bash ./preprocess_HMDB.sh $DATA_DIR/raw_data $DATA_DIR/processed_data\n",
    "\n",
    "        # download the split files and unrar\n",
    "        !wget -P $DATA_DIR --no-check-certificate http://serre-lab.clps.brown.edu/wp-content/uploads/2013/10/test_train_splits.rar\n",
    "        assert os.path.exists(f\"{DATA_DIR}/test_train_splits.rar\")\n",
    "        !mkdir -p $DATA_DIR/splits && unrar x -o+ $DATA_DIR/test_train_splits.rar $DATA_DIR/splits\n",
    "        assert os.path.exists(f\"{DATA_DIR}/splits\")\n",
    "        # run split_HMDB to generate training split\n",
    "        !if [ -d $DATA_DIR/train ]; then rm -rf $DATA_DIR/train $DATA_DIR/test; fi\n",
    "        !cd tao_toolkit_recipes/tao_action_recognition/data_generation/ && python3 ./split_dataset.py $DATA_DIR/processed_data $DATA_DIR/splits/testTrainMulti_7030_splits $DATA_DIR/train  $DATA_DIR/test\n",
    "        assert os.path.exists(f'{DATA_DIR}/train')\n",
    "        assert os.path.exists(f'{DATA_DIR}/test')\n",
    "\n",
    "        if os.path.exists(\"tao_toolkit_recipes\"):\n",
    "            shutil.rmtree(\"tao_toolkit_recipes\")\n",
    "\n",
    "        assert not os.path.exists(\"tao_toolkit_recipes\")\n",
    "\n",
    "    elif model_name == \"fpenet\":\n",
    "        !pip3 install numpy opencv-python\n",
    "        if model_type == \"80\":\n",
    "            output_json_path = os.path.join(os.environ['DATA_DIR'], 'data/afw/afw.json')\n",
    "        elif model_type == \"10\":\n",
    "            output_json_path = os.path.join(os.environ['DATA_DIR'], 'data/afw_10/afw_10.json')\n",
    "        !python3 ../dataset_prepare/fpenet/data_utils.py --afw_data_path $DATA_DIR/data/afw --output_json_path $output_json_path --afw_image_save_path $DATA_DIR/data/afw --num_key_points $model_type --container_root_path $USER_EXPERIMENT_DIR\n",
    "        assert os.path.exists(output_json_path)\n",
    "        with open(output_json_path, encoding='utf-8') as afw_json_file:\n",
    "            afw_json = json.load(afw_json_file)\n",
    "            assert afw_json\n",
    "\n",
    "    elif model_name == \"lprnet\":\n",
    "        character_file_link = \"https://api.ngc.nvidia.com/v2/models/nvidia/tao/lprnet/versions/trainable_v1.0/files/{}_lp_characters.txt\".format(model_type)\n",
    "        !wget -q -O $DATA_DIR/train/characters.txt $character_file_link\n",
    "        !cp $DATA_DIR/train/characters.txt $DATA_DIR/val/characters.txt\n",
    "        assert os.path.exists(f\"{DATA_DIR}/train/characters.txt\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/val/characters.txt\")\n",
    "\n",
    "    elif model_name == \"ml_recog\":\n",
    "        # crops images from detection set and form a classification set\n",
    "        # splits to reference/train/val/test set\n",
    "        !sudo apt-get update && sudo apt-get install gcc -y\n",
    "        !python3 -m pip install opencv-python numpy pycocotools tqdm\n",
    "        !python3 ../dataset_prepare/metric_learning_recognition/process_retail_product_checkout_dataset.py\n",
    "        assert os.path.exists(f\"{DATA_DIR}/metric_learning_recognition/retail-product-checkout-dataset_classification_demo\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/metric_learning_recognition/retail-product-checkout-dataset_classification_demo/known_classes\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/metric_learning_recognition/retail-product-checkout-dataset_classification_demo/unknown_classes\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/metric_learning_recognition/retail-product-checkout-dataset_classification_demo/known_classes/train\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/metric_learning_recognition/retail-product-checkout-dataset_classification_demo/unknown_classes/train\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/metric_learning_recognition/retail-product-checkout-dataset_classification_demo/known_classes/test\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/metric_learning_recognition/retail-product-checkout-dataset_classification_demo/unknown_classes/test\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/metric_learning_recognition/retail-product-checkout-dataset_classification_demo/known_classes/val\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/metric_learning_recognition/retail-product-checkout-dataset_classification_demo/unknown_classes/val\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/metric_learning_recognition/retail-product-checkout-dataset_classification_demo/known_classes/reference\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/metric_learning_recognition/retail-product-checkout-dataset_classification_demo/unknown_classes/reference\")\n",
    "\n",
    "    elif model_name == \"ocrnet\":\n",
    "        orig_train_gt_file=os.path.join(os.getenv(\"DATA_DIR\"), \"train\", \"gt.txt\")\n",
    "        processed_train_gt_file=os.path.join(os.getenv(\"DATA_DIR\"), \"train\", \"gt_new.txt\")\n",
    "        orig_test_gt_file=os.path.join(os.getenv(\"DATA_DIR\"), \"test\", \"Challenge4_Test_Task3_GT.txt\")\n",
    "        processed_test_gt_file=os.path.join(os.getenv(\"DATA_DIR\"), \"test\", \"gt_new.txt\")\n",
    "        !python3 ../dataset_prepare/ocrnet/preprocess_label.py $orig_train_gt_file $processed_train_gt_file\n",
    "        !python3 ../dataset_prepare/ocrnet/preprocess_label.py $orig_test_gt_file $processed_test_gt_file\n",
    "\n",
    "    elif model_name == \"pointpillars\":\n",
    "        !python3 -m pip install scikit-image numpy\n",
    "        !mkdir -p $DATA_DIR/train/lidar $DATA_DIR/train/label $DATA_DIR/val/lidar $DATA_DIR/val/label\n",
    "\n",
    "        # Convert labels from Camera coordinate system to LIDAR coordinate system, etc\n",
    "        !python3 ../dataset_prepare/pointpillars/gen_lidar_points.py -p $DATA_DIR/training/velodyne \\\n",
    "                                               -c $DATA_DIR/training/calib    \\\n",
    "                                               -i $DATA_DIR/training/image_2  \\\n",
    "                                               -o $DATA_DIR/train/lidar\n",
    "        assert os.listdir(f\"{DATA_DIR}/train/lidar\")\n",
    "        # Drop DontCare class\n",
    "        !python3 ../dataset_prepare/pointpillars/gen_lidar_labels.py -l $DATA_DIR/training/label_2 \\\n",
    "                                               -c $DATA_DIR/training/calib \\\n",
    "                                               -o $DATA_DIR/train/label\n",
    "        # train/val split\n",
    "        !python3 ../dataset_prepare/pointpillars/drop_class.py $DATA_DIR/train/label DontCare\n",
    "        assert os.listdir(f\"{DATA_DIR}/train/label\")\n",
    "        # Change the val set id's if you need a different set of validation images\n",
    "        !python3 ../dataset_prepare/pointpillars/kitti_split.py ../dataset_prepare/pointpillars/val.txt \\\n",
    "                                          $DATA_DIR/train/lidar \\\n",
    "                                          $DATA_DIR/train/label \\\n",
    "                                          $DATA_DIR/val/lidar \\\n",
    "                                          $DATA_DIR/val/label\n",
    "        assert os.listdir(f\"{DATA_DIR}/val/label\")\n",
    "        assert os.listdir(f\"{DATA_DIR}/val/lidar\")\n",
    "\n",
    "    elif model_name == \"pose_classification\" and model_type == \"kinetics\":\n",
    "        !pip3 install numpy\n",
    "        # select actions\n",
    "        !python3 ../dataset_prepare/pose_classification/select_subset_actions.py\n",
    "        assert os.path.exists(f\"{DATA_DIR}/{model_type}/train_data.npy\") and os.path.exists(f\"{DATA_DIR}/{model_type}/train_label.pkl\") and os.path.exists(f\"{DATA_DIR}/{model_type}/val_data.npy\") and os.path.exists(f\"{DATA_DIR}/{model_type}/val_label.pkl\")\n",
    "\n",
    "    elif model_name == \"re_identification\":\n",
    "        #100 is the number of samples to be present in the subset data - you can choose any number <= total samples in the dataset\n",
    "        !python3 ../dataset_prepare/re_identification/obtain_subset_data.py 100\n",
    "        assert os.path.exists(f\"{DATA_DIR}/market1501/sample_train\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/market1501/sample_test\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/market1501/sample_query\")\n",
    "\n",
    "    elif model_name == \"centerpose\":\n",
    "        # Select the training categories from: bike, book, bottle, camera, cereal_box, chair, laptop, shoe\n",
    "        # Please set the \"n\" to -1 if you want to run the whole dataset training.\n",
    "        testing_categories = 'bike'\n",
    "        !pip3 install numpy opencv-python tqdm scipy==1.9.2 tensorflow==2.14.0\n",
    "        !python3 ../dataset_prepare/centerpose/prepare_centerpose_dataset.py \\\n",
    "                                            -c $testing_categories \\\n",
    "                                            -n 100\n",
    "        assert os.path.exists(f\"{DATA_DIR}/{testing_categories}/train\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/{testing_categories}/test\")\n",
    "        assert os.path.exists(f\"{DATA_DIR}/{testing_categories}/val\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Update\n",
    "docker_env_vars = {} # Update any variables to be included while triggering Docker run-time like MLOPs variables \n",
    "dataset_information = {\"name\":\"Train dataset\",\n",
    "                       \"description\":\"My train dataset\",\n",
    "                       \"docker_env_vars\": docker_env_vars}\n",
    "data = json.dumps(dataset_information)\n",
    "\n",
    "endpoint = f\"{base_url}/datasets/{train_dataset_id}\"\n",
    "\n",
    "response = requests.patch(endpoint, data=data, headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "\n",
    "print(response)\n",
    "print(response.json())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Tar the datasets <a class=\"anchor\" id=\"head-1.3\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if model_name == \"action_recognition\":\n",
    "    !tar -C $DATA_DIR -czf $DATA_DIR/purpose_built_models_train.tar.gz train test\n",
    "elif model_name == \"bpnet\":\n",
    "    !tar -C $DATA_DIR -czf $DATA_DIR/purpose_built_models_train.tar.gz train2017 val2017 annotations bpnet_18joints.json  coco_spec.json  infer_spec.yaml\n",
    "elif model_name == \"fpenet\":\n",
    "    !tar -C $DATA_DIR -czf $DATA_DIR/purpose_built_models_train.tar.gz data data.json\n",
    "elif model_name == \"lprnet\":\n",
    "    !tar -C $DATA_DIR/train/ -czf $DATA_DIR/purpose_built_models_train.tar.gz image label characters.txt\n",
    "    !tar -C $DATA_DIR/val/ -czf $DATA_DIR/purpose_built_models_val.tar.gz image label characters.txt\n",
    "    !tar -C $DATA_DIR/val/ -czf $DATA_DIR/purpose_built_models_test.tar.gz image\n",
    "elif model_name == \"ml_recog\":\n",
    "    !tar -C $DATA_DIR -czf $DATA_DIR/purpose_built_models_train.tar.gz metric_learning_recognition/retail-product-checkout-dataset_classification_demo/\n",
    "elif model_name == \"ocdnet\":\n",
    "    !tar -C $DATA_DIR -czf $DATA_DIR/purpose_built_models_train.tar.gz train\n",
    "    !tar -C $DATA_DIR -czf $DATA_DIR/purpose_built_models_val.tar.gz test\n",
    "elif model_name == \"ocrnet\":\n",
    "    !tar -C $DATA_DIR -czf $DATA_DIR/purpose_built_models_train.tar.gz train character_list\n",
    "    !tar -C $DATA_DIR -czf $DATA_DIR/purpose_built_models_val.tar.gz test character_list\n",
    "elif model_name == \"optical_inspection\" or ds_format == \"visual_changenet_classify\":\n",
    "    !tar -C $DATA_DIR/train -czf $DATA_DIR/purpose_built_models_train.tar.gz images dataset.csv\n",
    "    !tar -C $DATA_DIR/val -czf $DATA_DIR/purpose_built_models_val.tar.gz images dataset.csv\n",
    "    !tar -C $DATA_DIR/test -czf $DATA_DIR/purpose_built_models_test.tar.gz images dataset.csv\n",
    "elif model_name == \"visual_changenet\" and ds_format == \"visual_changenet_segment\":\n",
    "    !tar -C $DATA_DIR/LEVIR-CD256 -czf $DATA_DIR/purpose_built_models_train.tar.gz A B list label\n",
    "elif model_name == \"pointpillars\":\n",
    "    !tar -C $DATA_DIR -czf $DATA_DIR/purpose_built_models_train.tar.gz train val\n",
    "elif model_name == \"pose_classification\":\n",
    "    !tar -C $DATA_DIR -czf $DATA_DIR/purpose_built_models_train.tar.gz $model_type\n",
    "elif model_name == \"re_identification\":\n",
    "    !tar -C $DATA_DIR/market1501 -czf $DATA_DIR/purpose_built_models_train.tar.gz sample_train sample_test sample_query\n",
    "elif model_name == 'centerpose':\n",
    "    !tar -C $DATA_DIR/{testing_categories} -czf $DATA_DIR/purpose_built_models_train.tar.gz train val test"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Upload train dataset <a class=\"anchor\" id=\"head-1.3\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Upload\n",
    "output_dir = os.path.join(os.path.dirname(os.path.abspath(train_dataset_path)), model_name, \"train\")\n",
    "split_tar_file(train_dataset_path, output_dir)\n",
    "for idx, tar_dataset_path in enumerate(os.listdir(output_dir)):\n",
    "    print(f\"Uploading {idx+1}/{len(os.listdir(output_dir))} tar split\")\n",
    "    files = [(\"file\",open(os.path.join(output_dir, tar_dataset_path),\"rb\"))]\n",
    "\n",
    "    endpoint = f\"{base_url}/datasets/{train_dataset_id}:upload\"\n",
    "\n",
    "    response = requests.post(endpoint, files=files, headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "    assert \"message\" in response.json().keys() and response.json()[\"message\"] == \"Server recieved file and upload process started\"\n",
    "\n",
    "    print(response)\n",
    "    print(response.json())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create and upload val dataset <a class=\"anchor\" id=\"head-1.3\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Create eval dataset\n",
    "if model_name in (\"lprnet\", \"ocdnet\", \"ocrnet\", \"optical_inspection\") or ds_format in (\"visual_changenet_classify\"):\n",
    "    data = json.dumps({\"type\":ds_type,\"format\":ds_format})\n",
    "\n",
    "    endpoint = f\"{base_url}/datasets\"\n",
    "\n",
    "    response = requests.post(endpoint,data=data,headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "\n",
    "    print(response)\n",
    "    print(response.json())\n",
    "    assert \"id\" in response.json().keys()\n",
    "    eval_dataset_id = response.json()[\"id\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Update\n",
    "if model_name in (\"lprnet\", \"ocdnet\", \"ocrnet\", \"optical_inspection\") or ds_format in (\"visual_changenet_classify\"):\n",
    "    docker_env_vars = {} # Update any variables to be included while triggering Docker run-time like MLOPs variables \n",
    "    dataset_information = {\"name\":\"Eval dataset\",\n",
    "                           \"description\":\"My eval dataset with OpenALPR\",\n",
    "                           \"docker_env_vars\": docker_env_vars}\n",
    "    data = json.dumps(dataset_information)\n",
    "\n",
    "    endpoint = f\"{base_url}/datasets/{eval_dataset_id}\"\n",
    "\n",
    "    response = requests.patch(endpoint, data=data, headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "\n",
    "    print(response)\n",
    "    print(response.json())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Upload\n",
    "if model_name in (\"lprnet\", \"ocdnet\", \"ocrnet\", \"optical_inspection\") or ds_format in (\"visual_changenet_classify\"):\n",
    "    output_dir = os.path.join(os.path.dirname(os.path.abspath(eval_dataset_path)), model_name, \"eval\")\n",
    "    split_tar_file(eval_dataset_path, output_dir)\n",
    "    for idx, tar_dataset_path in enumerate(os.listdir(output_dir)):\n",
    "        print(f\"Uploading {idx+1}/{len(os.listdir(output_dir))} tar split\")\n",
    "        files = [(\"file\",open(os.path.join(output_dir, tar_dataset_path),\"rb\"))]\n",
    "\n",
    "        endpoint = f\"{base_url}/datasets/{eval_dataset_id}:upload\"\n",
    "\n",
    "        response = requests.post(endpoint, files=files, headers=headers)\n",
    "        assert response.status_code in (200, 201)\n",
    "        assert \"message\" in response.json().keys() and response.json()[\"message\"] == \"Server recieved file and upload process started\"\n",
    "\n",
    "        print(response)\n",
    "        print(response.json())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create and upload test dataset <a class=\"anchor\" id=\"head-1.4\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Create testing dataset for inference\n",
    "if model_name in (\"lprnet\", \"optical_inspection\") or ds_format in (\"visual_changenet_classify\"):\n",
    "    if model_name == \"lprnet\":\n",
    "        ds_type = \"character_recognition\"\n",
    "        ds_format = \"raw\"\n",
    "    elif ds_format == \"visual_changenet_classify\": \n",
    "        ds_type = \"visual_changenet\"\n",
    "        ds_format = \"visual_changenet_classify\"\n",
    "    else:\n",
    "        ds_type = model_name\n",
    "        ds_format = \"default\"\n",
    "\n",
    "    data = json.dumps({\"type\":ds_type,\"format\":ds_format})\n",
    "\n",
    "    endpoint = f\"{base_url}/datasets\"\n",
    "\n",
    "    response = requests.post(endpoint,data=data, headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "\n",
    "    print(response)\n",
    "    print(response.json())\n",
    "    assert \"id\" in response.json().keys()\n",
    "    test_dataset_id = response.json()[\"id\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Upload\n",
    "if model_name in (\"lprnet\", \"optical_inspection\") or ds_format in (\"visual_changenet_classify\"):\n",
    "    output_dir = os.path.join(os.path.dirname(os.path.abspath(test_dataset_path)), model_name, \"test\")\n",
    "    split_tar_file(test_dataset_path, output_dir)\n",
    "    for idx, tar_dataset_path in enumerate(os.listdir(output_dir)):\n",
    "        print(f\"Uploading {idx+1}/{len(os.listdir(output_dir))} tar split\")\n",
    "        files = [(\"file\",open(os.path.join(output_dir, tar_dataset_path),\"rb\"))]\n",
    "\n",
    "        endpoint = f\"{base_url}/datasets/{test_dataset_id}:upload\"\n",
    "\n",
    "        response = requests.post(endpoint, files=files, headers=headers)\n",
    "        assert response.status_code in (200, 201)\n",
    "        assert \"message\" in response.json().keys() and response.json()[\"message\"] == \"Server recieved file and upload process started\"\n",
    "\n",
    "        print(response)\n",
    "        print(response.json())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### List the created datasets <a class=\"anchor\" id=\"head-2\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "endpoint = f\"{base_url}/datasets\"\n",
    "\n",
    "response = requests.get(endpoint, headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "\n",
    "print(response)\n",
    "# print(response.json()) ## Uncomment for verbose list output\n",
    "print(\"id\\t\\t\\t\\t\\t type\\t\\t\\t format\\t\\t name\")\n",
    "for rsp in response.json():\n",
    "    rsp_keys = rsp.keys()\n",
    "    assert \"id\" in rsp_keys\n",
    "    assert \"type\" in rsp_keys\n",
    "    assert \"format\" in rsp_keys\n",
    "    assert \"name\" in rsp_keys\n",
    "    print(rsp[\"id\"],\"\\t\",rsp[\"type\"],\"\\t\",rsp[\"format\"],\"\\t\\t\",rsp[\"name\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train Dataset convert Action <a class=\"anchor\" id=\"head-3\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "convert_action = \"dataset_convert\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "if model_name in (\"bpnet\", \"fpenet\", \"ocrnet\", \"pointpillars\"):\n",
    "    # Get default spec schema\n",
    "    endpoint = f\"{base_url}/datasets/{train_dataset_id}/specs/{convert_action}/schema\"\n",
    "\n",
    "    response = requests.get(endpoint, headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "\n",
    "    print(response)\n",
    "    # print(response.json()) ## Uncomment for verbose schema\n",
    "\n",
    "    assert \"default\" in response.json().keys()\n",
    "    train_ds_convert_specs = response.json()[\"default\"]\n",
    "\n",
    "    print(json.dumps(train_ds_convert_specs, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Apply changes to specs dictionary if necessary\n",
    "if model_name in (\"bpnet\", \"fpenet\", \"ocrnet\", \"pointpillars\"):\n",
    "    if model_name == \"bpnet\":\n",
    "        train_ds_convert_specs[\"mode\"] = \"train\"\n",
    "    elif model_name == \"fpenet\":\n",
    "        train_ds_convert_specs[\"num_keypoints\"] = int(model_type)\n",
    "    print(json.dumps(train_ds_convert_specs, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Run action\n",
    "if model_name in (\"bpnet\", \"fpenet\", \"ocrnet\", \"pointpillars\"):\n",
    "    parent = None\n",
    "    action = convert_action\n",
    "    data = json.dumps({\"parent_job_id\":parent,\"action\":action, \"specs\":train_ds_convert_specs})\n",
    "\n",
    "    endpoint = f\"{base_url}/datasets/{train_dataset_id}/jobs\"\n",
    "\n",
    "    response = requests.post(endpoint, data=data, headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "    assert response.json()\n",
    "    print(response)\n",
    "    print(response.json())\n",
    "\n",
    "    train_ds_convert_id = response.json()\n",
    "    job_map[\"train_dataset_convert_\"+model_name] = train_ds_convert_id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Monitor job status by repeatedly running this cell\n",
    "if model_name in (\"bpnet\", \"fpenet\", \"ocrnet\", \"pointpillars\"):\n",
    "    job_id = train_ds_convert_id\n",
    "    endpoint = f\"{base_url}/datasets/{train_dataset_id}/jobs/{job_id}\"\n",
    "\n",
    "    while True:\n",
    "        clear_output(wait=True) \n",
    "        response = requests.get(endpoint, headers=headers)\n",
    "        assert response.status_code in (200, 201)\n",
    "        print(response)\n",
    "        print(response.json())\n",
    "        assert \"status\" in response.json().keys() and response.json().get(\"status\") != \"Error\"\n",
    "        if response.json().get(\"status\") in [\"Done\",\"Error\", \"Canceled\"] or response.status_code not in (200,201):\n",
    "            break\n",
    "        time.sleep(15)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Eval Dataset convert Action <a class=\"anchor\" id=\"head-3\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "if model_name in (\"bpnet\", \"ocrnet\"):\n",
    "    # Get default spec schema\n",
    "    if model_name == \"bpnet\":\n",
    "        endpoint = f\"{base_url}/datasets/{train_dataset_id}/specs/{convert_action}/schema\"\n",
    "    else:\n",
    "        endpoint = f\"{base_url}/datasets/{eval_dataset_id}/specs/{convert_action}/schema\"\n",
    "\n",
    "    response = requests.get(endpoint, headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "\n",
    "    print(response)\n",
    "    # print(response.json()) ## Uncomment for verbose schema\n",
    "\n",
    "    assert \"default\" in response.json().keys()\n",
    "    eval_ds_convert_specs = response.json()[\"default\"]\n",
    "\n",
    "    print(json.dumps(eval_ds_convert_specs, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Apply changes to specs dictionary if necessary\n",
    "if model_name in (\"bpnet\", \"ocrnet\"):\n",
    "    if model_name == \"bpnet\":\n",
    "        eval_ds_convert_specs[\"mode\"] = \"test\"\n",
    "    print(json.dumps(eval_ds_convert_specs, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Run action\n",
    "if model_name in (\"bpnet\", \"ocrnet\"):\n",
    "    parent = job_map[\"train_dataset_convert_\"+model_name]\n",
    "    action = convert_action\n",
    "    data = json.dumps({\"parent_job_id\":parent,\"action\":action,\"specs\":eval_ds_convert_specs})\n",
    "\n",
    "    if model_name == \"bpnet\":\n",
    "        endpoint = f\"{base_url}/datasets/{train_dataset_id}/jobs\"\n",
    "    else:\n",
    "        endpoint = f\"{base_url}/datasets/{eval_dataset_id}/jobs\"\n",
    "    \n",
    "    response = requests.post(endpoint, data=data, headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "    assert response.json()\n",
    "\n",
    "    print(response)\n",
    "    print(response.json())\n",
    "\n",
    "    eval_ds_convert_id = response.json()\n",
    "    job_map[\"eval_dataset_convert_\"+model_name] = eval_ds_convert_id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Monitor job status by repeatedly running this cell\n",
    "if model_name in (\"bpnet\", \"ocrnet\"):\n",
    "    job_id = eval_ds_convert_id\n",
    "    if model_name == \"bpnet\":\n",
    "        endpoint = f\"{base_url}/datasets/{train_dataset_id}/jobs/{job_id}\"\n",
    "    else:\n",
    "        endpoint = f\"{base_url}/datasets/{eval_dataset_id}/jobs/{job_id}\"\n",
    "\n",
    "    while True:\n",
    "        clear_output(wait=True) \n",
    "        response = requests.get(endpoint, headers=headers)\n",
    "        assert response.status_code in (200, 201)\n",
    "        print(response)\n",
    "        print(response.json())\n",
    "        assert \"status\" in response.json().keys() and response.json().get(\"status\") != \"Error\"\n",
    "        if response.json().get(\"status\") in [\"Done\",\"Error\", \"Canceled\"] or response.status_code not in (200,201):\n",
    "            break\n",
    "        time.sleep(15)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create an experiment <a class=\"anchor\" id=\"head-4\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if model_name in (\"action_recognition\", \"pose_classification\", \"ml_recog\", \"ocrnet\", \"ocdnet\", \"optical_inspection\", \"re_identification\"):\n",
    "    encode_key = \"nvidia_tao\"\n",
    "elif model_name == \"pointpillars\":\n",
    "    encode_key = \"tlt_encode\"\n",
    "else:\n",
    "    encode_key = \"nvidia_tlt\"\n",
    "\n",
    "checkpoint_choose_method = \"best_model\"\n",
    "data = json.dumps({\"network_arch\":model_name,\"encryption_key\":encode_key,\"checkpoint_choose_method\":checkpoint_choose_method})\n",
    "\n",
    "endpoint = f\"{base_url}/experiments\"\n",
    "response = requests.post(endpoint,data=data,headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "\n",
    "print(response)\n",
    "print(response.json())\n",
    "assert \"id\" in response.json().keys()\n",
    "experiment_id = response.json()[\"id\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### List experiments <a class=\"anchor\" id=\"head-5\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "endpoint = f\"{base_url}/experiments\"\n",
    "params = {\"network_arch\": model_name}\n",
    "response = requests.get(endpoint, params=params, headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "\n",
    "print(response)\n",
    "# print(response.json()) ## Uncomment for verbose list output\n",
    "print(\"model id\\t\\t\\t     network architecture\")\n",
    "for rsp in response.json():\n",
    "    rsp_keys = rsp.keys()\n",
    "    assert \"id\" in rsp_keys and \"network_arch\" in rsp_keys\n",
    "    print(rsp[\"name\"], rsp[\"id\"], rsp[\"network_arch\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Assign train, eval datasets <a class=\"anchor\" id=\"head-6\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "docker_env_vars = {} # Update any variables to be included while triggering Docker run-time like MLOPs variables \n",
    "dataset_information = {}\n",
    "dataset_information[\"train_datasets\"] = [train_dataset_id]\n",
    "if model_name in (\"bpnet\",\"fpenet\",\"lprnet\",\"ml_recog\",\"ocdnet\",\"ocrnet\"):\n",
    "    dataset_information[\"calibration_dataset\"] = train_dataset_id\n",
    "if model_name in (\"lprnet\", \"ocdnet\", \"ocrnet\", \"optical_inspection\"):\n",
    "    dataset_information[\"eval_dataset\"] = eval_dataset_id\n",
    "if model_name in (\"lprnet\", \"optical_inspection\"):\n",
    "    dataset_information[\"inference_dataset\"] = test_dataset_id\n",
    "if model_name in (\"centerpose\"):\n",
    "    dataset_information[\"eval_dataset\"] = train_dataset_id\n",
    "    dataset_information[\"inference_dataset\"] = train_dataset_id\n",
    "if model_name in (\"visual_changenet\") and ds_format in (\"visual_changenet_classify\"):\n",
    "    dataset_information[\"eval_dataset\"] = eval_dataset_id\n",
    "    dataset_information[\"inference_dataset\"] = test_dataset_id\n",
    "\n",
    "dataset_information[\"docker_env_vars\"] = docker_env_vars\n",
    "\n",
    "data = json.dumps(dataset_information)\n",
    "\n",
    "endpoint = f\"{base_url}/experiments/{experiment_id}\"\n",
    "\n",
    "response = requests.patch(endpoint, data=data, headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "\n",
    "print(response)\n",
    "print(response.json())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Assign PTM <a class=\"anchor\" id=\"head-7\"></a>\n",
    "\n",
    "Search for PTM on NGC for the Purpose built model chosen"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# List all pretrained models for the chosen network architecture\n",
    "endpoint = f\"{base_url}/experiments\"\n",
    "params = {\"network_arch\": model_name}\n",
    "response = requests.get(endpoint, params=params, headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "\n",
    "response_json = response.json()\n",
    "\n",
    "# Search for ptm with given ngc path\n",
    "for rsp in response_json:\n",
    "    rsp_keys = rsp.keys()\n",
    "    if \"encryption_key\" not in rsp.keys():\n",
    "        assert \"name\" in rsp_keys and \"version\" in rsp_keys and \"ngc_path\" in rsp_keys and \"additional_id_info\" in rsp_keys\n",
    "        print(f'PTM Name: {rsp[\"name\"]}; PTM version: {rsp[\"version\"]}; NGC PATH: {rsp[\"ngc_path\"]}; Additional info: {rsp[\"additional_id_info\"]}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Assigning pretrained models to different purpose built models versions\n",
    "# From the output of previous cell make the appropriate changes to this map if you want to change the default PTM backbone.\n",
    "# Changing the default backbone here requires changing default spec/config during train/eval etc like for example\n",
    "# If you are changing the ptm to resnet34, then you have to modify the config key num_layers if it exists to 34 manually\n",
    "visual_changenet_ptm = \"visual_changenet_segmentation_levircd:visual_changenet_levircd_trainable_v1.0\" # For segmentation\n",
    "if model_name == 'visual_changenet' and ds_format == 'visual_changenet_classify':\n",
    "    visual_changenet_ptm = \"visual_changenet_classification:visual_changenet_nvpcb_trainable_v1.0\"\n",
    "pretrained_map = {\"action_recognition\":\"actionrecognitionnet:trainable_v1.0\",\n",
    "                  \"bpnet\" : \"bodyposenet:trainable_v1.0\",\n",
    "                  \"fpenet\" : \"fpenet:trainable_v1.0\",\n",
    "                  \"lprnet\": \"lprnet:trainable_v1.0\",\n",
    "                  \"ml_recog\": \"retail_object_recognition:trainable_v1.0\",\n",
    "                  \"ocdnet\": \"ocdnet:trainable_resnet18_v1.0\",\n",
    "                  \"ocrnet\": \"ocrnet:trainable_v1.0\",\n",
    "                  \"optical_inspection\": \"optical_inspection:trainable_v1.0\",\n",
    "                  \"pointpillars\":\"pointpillarnet:trainable_v1.0\",\n",
    "                  \"pose_classification\":\"poseclassificationnet:trainable_v1.0\",\n",
    "                  \"re_identification\":\"reidentificationnet:trainable_v1.1\",\n",
    "                  \"visual_changenet\":visual_changenet_ptm,\n",
    "                  \"centerpose\": \"pretrained_fan_classification_nvimagenet:fan_small_hybrid_nvimagenet\"}\n",
    "if model_name == \"action_recognition\":\n",
    "    if model_type == \"of\":\n",
    "        pretrained_map[\"action_recognition\"] = \"actionrecognitionnet:trainable_v2.0\"\n",
    "    elif model_type == \"joint\":\n",
    "        pretrained_map[\"action_recognition\"] = \"actionrecognitionnet:trainable_v1.0,actionrecognitionnet:trainable_v2.0\"\n",
    "        \n",
    "no_ptm_models = set([])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if model_name not in no_ptm_models:\n",
    "    # Get pretrained model\n",
    "    endpoint = f\"{base_url}/experiments\"\n",
    "    params = {\"network_arch\": model_name}\n",
    "    response = requests.get(endpoint, params=params, headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "\n",
    "    response_json = response.json()\n",
    "    ptm_model_names = pretrained_map[model_name].split(\",\")\n",
    "    ptm = []\n",
    "\n",
    "    # Search for ptm with given ngc path\n",
    "    for ptm_model_name in ptm_model_names:\n",
    "        ptm_id = None\n",
    "        for rsp in response_json:\n",
    "            rsp_keys = rsp.keys()\n",
    "            assert \"ngc_path\" in rsp_keys\n",
    "            if rsp[\"ngc_path\"].endswith(ptm_model_name):\n",
    "                additional_id_info = []\n",
    "                if rsp[\"additional_id_info\"]:\n",
    "                    assert \"additional_id_info\" in rsp_keys\n",
    "                    additional_id_info = rsp[\"additional_id_info\"].split(\",\")\n",
    "                if (len(additional_id_info) == 0) or \\\n",
    "                    (model_name == \"lprnet\" and len(additional_id_info) == 1 and additional_id_info[0] == model_type) or \\\n",
    "                    (model_name == \"action_recognition\" and len(additional_id_info) == 1 and additional_id_info[0] == model_input_type) or \\\n",
    "                    (model_name == \"action_recognition\" and len(additional_id_info) == 2 and additional_id_info[0] == platform and additional_id_info[1] == model_input_type):\n",
    "                    assert \"id\" in rsp_keys\n",
    "                    ptm_id = rsp[\"id\"]\n",
    "                    print(\"Metadata for model with requested NGC Path\")\n",
    "                    print(rsp)\n",
    "                    break\n",
    "        ptm.append(ptm_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if model_name not in no_ptm_models:\n",
    "    ptm_information = {\"base_experiment\":ptm}\n",
    "    data = json.dumps(ptm_information)\n",
    "    endpoint = f\"{base_url}/experiments/{experiment_id}\"\n",
    "\n",
    "    response = requests.patch(endpoint, data=data, headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "\n",
    "    print(response)\n",
    "    print(response.json())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### View hyperparameters that are enabled for AutoML by default <a class=\"anchor\" id=\"head-8\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if automl_enabled:\n",
    "    # Get default spec schema\n",
    "    endpoint = f\"{base_url}/experiments/{experiment_id}/specs/train/schema\"\n",
    "    response = requests.get(endpoint, headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "    assert \"automl_default_parameters\" in response.json().keys()\n",
    "    automl_specs = response.json()[\"automl_default_parameters\"]\n",
    "    print(json.dumps(automl_specs, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Actions <a class=\"anchor\" id=\"head-10\"></a>\n",
    "\n",
    "For all actions:\n",
    "1. Get default spec schema and derive the default values\n",
    "2. Modify defaults if needed\n",
    "3. Post spec dictionary to the service\n",
    "4. Run model action\n",
    "5. Monitor job using retrieve\n",
    "6. Download results using job download endpoint (if needed)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train <a class=\"anchor\" id=\"head-11\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Set AutoML related configurations <a class=\"anchor\" id=\"head-9\"></a>\n",
    "Refer to these hyper-links to see the parameters supported by each network and add more parameters if necessary in addition to the default automl enabled parameters:\n",
    "\n",
    "[ActionRecognitionNet](https://github.com/NVIDIA/tao_front_end_services/tree/main/api/specs_utils/specs/action_recognition/action_recognition%20-%20train.csv), \n",
    "[BPNET](https://github.com/NVIDIA/tao_front_end_services/tree/main/api/specs_utils/specs/bpet/bpnet%20-%20train.csv), \n",
    "[FPENET](https://github.com/NVIDIA/tao_front_end_services/tree/main/api/specs_utils/specs/fpenet/fpenet%20-%20train.csv), \n",
    "[LPRNET](https://github.com/NVIDIA/tao_front_end_services/tree/main/api/specs_utils/specs/lprnet/lprnet%20-%20train.csv), \n",
    "[MetricLearningRecognition](https://github.com/NVIDIA/tao_front_end_services/tree/main/api/specs_utils/specs/ml_recog/ml_recog%20-%20train.csv), \n",
    "[OCDNET](https://github.com/NVIDIA/tao_front_end_services/tree/main/api/specs_utils/specs/ocdnet/ocdnet%20-%20train.csv), \n",
    "[OCRNET](https://github.com/NVIDIA/tao_front_end_services/tree/main/api/specs_utils/specs/ocrnet/ocrnet%20-%20train.csv), \n",
    "[OpticalInspection](https://github.com/NVIDIA/tao_front_end_services/tree/main/api/specs_utils/specs/optical_inspection/optical_inspection%20-%20train.csv), \n",
    "[Pointpillars](https://github.com/NVIDIA/tao_front_end_services/tree/main/api/specs_utils/specs/pointpillars/pointpillars%20-%20train.csv), \n",
    "[PoseClassificationNet](https://github.com/NVIDIA/tao_front_end_services/tree/main/api/specs_utils/specs/pose_classification/pose_classification%20-%20train.csv), \n",
    "[ReIdentificationNet](https://github.com/NVIDIA/tao_front_end_services/tree/main/api/specs_utils/specs/re_identification/re_identification%20-%20train.csv)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if automl_enabled:\n",
    "    # Choose any metric that is present in the kpi dictionary present in the model's status.json. \n",
    "    # Example status.json for each model can be found in the respective section in NVIDIA TAO DOCS here: https://docs.nvidia.com/tao/tao-toolkit/text/model_zoo/cv_models/index.html\n",
    "    metric = \"kpi\"\n",
    "\n",
    "    additional_automl_parameters = [] #Refer to parameter list mentioned in the above links and add any extra parameter in addition to the default enabled ones\n",
    "    remove_default_automl_parameters = [] #Remove any hyperparameters that are enabled by default for AutoML\n",
    "\n",
    "    automl_information = {\"automl_enabled\":True,\n",
    "                          \"automl_algorithm\":automl_algorithm,\n",
    "                          \"metric\":metric,\n",
    "                          \"automl_max_recommendations\": 20, # Only for bayesian\n",
    "                          \"automl_R\": 27, # Only for hyperband\n",
    "                          \"automl_nu\": 3, # Only for hyperband\n",
    "                          \"epoch_multiplier\": 1, # Only for hyperband\n",
    "                          # Enable this if you want to add parameters to automl_add_hyperparameters below that are disabled by TAO in the automl_enabled column of the spec csv.\n",
    "                          # Warning: The parameters that are disabled are not tested by TAO, so there might be unexpected behaviour in overriding this\n",
    "                          \"override_automl_disabled_params\": False,\n",
    "                          \"automl_add_hyperparameters\":str(additional_automl_parameters),\n",
    "                          \"automl_remove_hyperparameters\":str(remove_default_automl_parameters)\n",
    "                        }\n",
    "    data = json.dumps(automl_information)\n",
    "\n",
    "    endpoint = f\"{base_url}/experiments/{experiment_id}\"\n",
    "\n",
    "    response = requests.patch(endpoint, data=data, headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "\n",
    "    print(response)\n",
    "    print(json.dumps(response.json(), sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Get default spec schema\n",
    "endpoint = f\"{base_url}/experiments/{experiment_id}/specs/train/schema\"\n",
    "\n",
    "response = requests.get(endpoint, headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "\n",
    "print(response)\n",
    "print(response.json()) ## Uncomment for verbose schema\n",
    "assert \"default\" in response.json().keys()\n",
    "train_specs = response.json()[\"default\"]\n",
    "print(json.dumps(train_specs, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Apply changes for any of the parameters listed in the previous cell as required\n",
    "if model_name == \"action_recognition\":\n",
    "    train_specs[\"model\"][\"model_type\"] = model_type\n",
    "    train_specs[\"model\"][\"input_type\"] = model_input_type\n",
    "    train_specs[\"train\"][\"num_epochs\"] = 20\n",
    "    train_specs[\"train\"][\"num_gpus\"] = 1\n",
    "elif model_name == \"bpnet\":\n",
    "    train_specs[\"num_epoch\"] = 20\n",
    "    train_specs[\"checkpoint_n_epoch\"] = 5\n",
    "    train_specs[\"validation_every_n_epoch\"] = 5\n",
    "    train_specs[\"finetuning_config\"][\"checkpoint_path\"] = None\n",
    "    train_specs[\"gpus\"] = 1\n",
    "elif model_name == \"centerpose\":\n",
    "    train_specs[\"train\"][\"num_epochs\"] = 20 # Please set it to 140 if you want to run the whole training.\n",
    "    train_specs[\"train\"][\"validation_interval\"] = 10\n",
    "    train_specs[\"train\"][\"checkpoint_interval\"] = 10\n",
    "    train_specs[\"train\"][\"num_gpus\"] = 1\n",
    "    train_specs[\"dataset\"][\"category\"] = testing_categories\n",
    "elif model_name == \"fpenet\":\n",
    "    train_specs[\"num_epoch\"] = 10\n",
    "    train_specs[\"checkpoint_n_epoch\"] = 5\n",
    "    train_specs[\"dataloader\"][\"dataset_info\"][\"root_path\"] = None\n",
    "    train_specs[\"num_keypoints\"] = int(model_type)\n",
    "    train_specs[\"dataloader\"][\"num_keypoints\"] = int(model_type)\n",
    "    train_specs[\"gpus\"] = 1\n",
    "elif model_name == \"lprnet\":\n",
    "    train_specs[\"training_config\"][\"num_epochs\"] = 24\n",
    "    train_specs[\"gpus\"] = 1\n",
    "elif model_name == \"ml_recog\":\n",
    "    train_specs[\"train\"][\"num_epochs\"] = 30\n",
    "    train_specs[\"train\"][\"gpu_ids\"] = [0]\n",
    "    train_specs[\"train\"][\"checkpoint_interval\"] = 5\n",
    "elif model_name == \"ocdnet\":\n",
    "    train_specs[\"train\"][\"num_epochs\"] = 30\n",
    "    train_specs[\"train\"][\"checkpoint_interval\"] = 5\n",
    "    train_specs[\"train\"][\"validation_interval\"] = 5\n",
    "    train_specs[\"train\"][\"gpu_id\"] = [0]\n",
    "    train_specs[\"num_gpus\"] = 1\n",
    "elif model_name == \"ocrnet\":\n",
    "    train_specs[\"train\"][\"num_epochs\"] = 20\n",
    "    train_specs[\"train\"][\"checkpoint_interval\"] = 5\n",
    "    train_specs[\"train\"][\"validation_interval\"] = 5\n",
    "    train_specs[\"train\"][\"num_gpus\"] = 1\n",
    "elif model_name == \"optical_inspection\":\n",
    "    train_specs[\"train\"][\"num_epochs\"] = 20\n",
    "    train_specs[\"train\"][\"checkpoint_interval\"] = 5\n",
    "    train_specs[\"train\"][\"validation_interval\"] = 5\n",
    "    train_specs[\"train\"][\"gpu_ids\"] = [0]\n",
    "elif model_name == \"pose_classification\":\n",
    "    train_specs[\"train\"][\"num_epochs\"] = 50\n",
    "    train_specs[\"train\"][\"gpu_ids\"] = [0]\n",
    "    train_specs[\"train\"][\"num_gpus\"] = 1\n",
    "    if model_type == \"nvidia\":\n",
    "        train_specs[\"dataset\"][\"num_classes\"] = 6\n",
    "        train_specs[\"model\"][\"graph_layout\"] = \"nvidia\"\n",
    "    elif model_type == \"kinetics\":\n",
    "        train_specs[\"dataset\"][\"num_classes\"] = 5\n",
    "        train_specs[\"model\"][\"graph_layout\"] = \"openpose\"\n",
    "elif model_name == \"pointpillars\":\n",
    "    train_specs[\"train\"][\"num_epochs\"] = 80\n",
    "    train_specs[\"gpus\"] = 1\n",
    "elif model_name == \"re_identification\":\n",
    "    train_specs[\"train\"][\"num_epochs\"] = 120\n",
    "    train_specs[\"train\"][\"gpu_ids\"] = [0]\n",
    "    train_specs[\"train\"][\"num_gpus\"] = 1\n",
    "    train_specs[\"dataset\"][\"num_classes\"] = 100 #The number set in obtain_subset script\n",
    "    train_specs[\"dataset\"][\"num_workers\"] = 4 #Modify the num_workers according to your hardware setup\n",
    "    train_specs[\"dataset\"][\"batch_size\"] = 16 #Modify the batch_size according to your hardware setup\n",
    "elif model_name == \"visual_changenet\":\n",
    "    train_specs[\"train\"][\"num_epochs\"] = 30 \n",
    "    train_specs[\"train\"][\"checkpoint_interval\"] = 2\n",
    "    train_specs[\"train\"][\"val_interval\"] = 5\n",
    "    train_specs[\"num_gpus\"] = 1\n",
    "    if ds_format == \"visual_changenet_segment\":\n",
    "        train_specs[\"task\"] = 'segment'\n",
    "    elif ds_format == \"visual_changenet_classify\":\n",
    "        train_specs[\"task\"] = 'classify'\n",
    "print(json.dumps(train_specs, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Run action\n",
    "parent = job_map.get(\"eval_dataset_convert_\"+model_name, job_map.get(\"train_dataset_convert_\"+model_name, None))\n",
    "parent_id = train_dataset_id\n",
    "if model_name == \"ocrnet\": # Only model with eval dataset convert on eval dataset\n",
    "    parent_id = eval_dataset_id\n",
    "action = \"train\"\n",
    "data = json.dumps({\"parent_job_id\":parent,\"action\":action,\"specs\":train_specs})\n",
    "endpoint = f\"{base_url}/experiments/{experiment_id}/jobs\"\n",
    "\n",
    "response = requests.post(endpoint, data=data, headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "assert response.json()\n",
    "\n",
    "print(response)\n",
    "print(response.json())\n",
    "\n",
    "if model_name == \"visual_changenet\":\n",
    "    job_map[\"train_\" + ds_format] = response.json()\n",
    "else:\n",
    "    job_map[\"train_\" + model_name] = response.json()\n",
    "print(job_map)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Monitor job status by repeatedly running this cell\n",
    "# For automl: Training times for different models benchmarked on 1 GPU V100 machine can be found here: https://docs.nvidia.com/tao/tao-toolkit/text/automl/automl.html#results-of-automl-experiments\n",
    "\n",
    "if model_name == \"visual_changenet\":\n",
    "    job_id = job_map[\"train_\" + ds_format]\n",
    "else:\n",
    "    job_id = job_map[\"train_\" + model_name]\n",
    "endpoint = f\"{base_url}/experiments/{experiment_id}/jobs/{job_id}\"\n",
    "\n",
    "while True:\n",
    "    clear_output(wait=True)\n",
    "    response = requests.get(endpoint, headers=headers)\n",
    "    if \"error_desc\" in response.json().keys() and response.json()[\"error_desc\"] in (\"Job trying to retrieve not found\", \"No AutoML run found\"):\n",
    "        print(\"Job is being created\")\n",
    "        time.sleep(5)\n",
    "        continue\n",
    "    assert response.status_code in (200, 201)\n",
    "    print(response)\n",
    "    print(json.dumps(response.json(), sort_keys=True, indent=4))\n",
    "    assert \"status\" in response.json().keys() and response.json().get(\"status\") != \"Error\"\n",
    "    if response.json().get(\"status\") in [\"Done\",\"Error\", \"Canceled\"] or response.status_code not in (200,201):\n",
    "        break\n",
    "    time.sleep(15)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "## To Stop an AutoML JOB\n",
    "#    1. Stop the 'Monitor job status by repeatedly running this cell' cell (the cell right before this cell) manually\n",
    "#    2. Uncomment the snippet in the next cell and run the cell"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# if automl_enabled:\n",
    "#     if model_name == \"visual_changenet\":\n",
    "#          job_id = job_map[\"train_\" + ds_format]\n",
    "#     else:\n",
    "#         job_id = job_map[\"train_\" + model_name]\n",
    "\n",
    "#     endpoint = f\"{base_url}/experiments/{experiment_id}/jobs/{job_id}:cancel\"\n",
    "\n",
    "#     response = requests.post(endpoint, headers=headers)\n",
    "#     assert response.status_code in (200, 201)\n",
    "\n",
    "#     print(response)\n",
    "#     print(response.json())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "## Resume AutoML"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Uncomment the below snippet if you want to resume an already stopped AutoML job and then run the 'Monitor job status by repeatedly running this cell' cell above (4th cell above from this cell)\n",
    "# if automl_enabled:\n",
    "#     if model_name == \"visual_changenet\":\n",
    "#          job_id = job_map[\"train_\" + ds_format]\n",
    "#     else:\n",
    "#         job_id = job_map[\"train_\" + model_name]\n",
    "#     endpoint = f\"{base_url}/experiments/{experiment_id}/jobs/{job_id}:resume\"\n",
    "\n",
    "#     data = json.dumps({\"parent_job_id\":parent,\"specs\":train_specs})\n",
    "#     response = requests.post(endpoint, data=data, headers=headers)\n",
    "#     assert response.status_code in (200, 201)\n",
    "\n",
    "#     print(response)\n",
    "#     print(response.json())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Download train job artifacts <a class=\"anchor\" id=\"head-12\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Example to list the files of the executed train job\n",
    "if model_name == \"visual_changenet\":\n",
    "    job_id = job_map[\"train_\" + ds_format]\n",
    "else:\n",
    "    job_id = job_map[\"train_\" + model_name]\n",
    "endpoint = f'{base_url}/experiments/{experiment_id}/jobs/{job_id}:list_files'\n",
    "\n",
    "response = requests.get(endpoint, headers=headers)\n",
    "print(json.dumps(response.json(), sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "## Patch the model with proper metric before training to run this cell; By default loss is used, but some models dont log the parameter under the name 'loss'\n",
    "\n",
    "# # Download selective job contents once the above job shows \"Done\" status\n",
    "# # Example to download selective files of train job (Note: will take time)\n",
    "# endpoint = f'{base_url}/experiments/{experiment_id}/jobs/{job_id}:download_selective_files'\n",
    "\n",
    "# file_lists = [] # Choose file names from the previous cell where all the files for this job were listed\n",
    "# best_model = False # Enable this to download the checkpoint of the best performing model w.r.t to the metric chosen before starting training\n",
    "# latest_model = True # Enable this to download the latest checkpoint of the training job; Disable best_model to use latest_model\n",
    "\n",
    "# params = {\"file_lists\": file_lists, \"best_model\": best_model, \"latest_model\": latest_model}\n",
    "\n",
    "# # Save\n",
    "# temptar = f'{job_id}.tar.gz'\n",
    "# with requests.get(endpoint, headers=headers, params=params, stream=True) as r:\n",
    "#     r.raise_for_status()\n",
    "#     with open(temptar, 'wb') as f:\n",
    "#         for chunk in r.iter_content(chunk_size=8192):\n",
    "#             f.write(chunk)\n",
    "\n",
    "# print(\"Untarring\")\n",
    "# # Untar to destination\n",
    "# tar_command = f'tar -xvf {temptar} -C {workdir}/'\n",
    "# os.system(tar_command)\n",
    "# os.remove(temptar)\n",
    "# print(f\"Results at {workdir}/{job_id}\")\n",
    "# model_downloaded_path = f\"{workdir}/{job_id}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Downloading train job takes a longer time, uncomment this cell if you want to still proceed\n",
    "if download_jobs:\n",
    "    endpoint = f\"{base_url}/experiments/{experiment_id}/jobs/{job_id}\"\n",
    "    response = requests.get(endpoint, headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "    expected_file_size = response.json().get(\"job_tar_stats\", {}).get(\"file_size\")\n",
    "    print(\"expected_file_size: \", expected_file_size)\n",
    "\n",
    "    !python3 -m pip install tqdm\n",
    "    from tqdm import tqdm\n",
    "\n",
    "    endpoint = f'{base_url}/experiments/{experiment_id}/jobs/{job_id}:download'\n",
    "    temptar = f'{job_id}.tar.gz'\n",
    "\n",
    "    with tqdm(total=expected_file_size, unit='B', unit_scale=True) as progress_bar:\n",
    "        while True:\n",
    "            # Check if the file already exists\n",
    "            headers_download_job = dict(headers)\n",
    "            if os.path.exists(temptar):\n",
    "                # Get the current file size\n",
    "                file_size = os.path.getsize(temptar)\n",
    "                print(f\"File size of dowloaded content until now is {file_size}\")\n",
    "\n",
    "                # If the file size matches the expected size, break out of the loop\n",
    "                if file_size >= (expected_file_size-1):\n",
    "                    print(\"Download completed successfully.\")\n",
    "                    print(\"Untarring\")\n",
    "                    # Untar to destination\n",
    "                    tar_command = f'tar -xf {temptar} -C {workdir}/'\n",
    "                    os.system(tar_command)\n",
    "                    os.remove(temptar)\n",
    "                    print(f\"Results at {workdir}/{job_id}\")\n",
    "                    model_downloaded_path = f\"{workdir}/{job_id}\"\n",
    "                    break\n",
    "\n",
    "                # Set the headers to resume the download from where it left off\n",
    "                headers_download_job['Range'] = f'bytes={file_size}-'\n",
    "            # Open the file for writing in binary mode\n",
    "            with open(temptar, 'ab') as f:\n",
    "                try:\n",
    "                    response = requests.get(endpoint, headers=headers_download_job, stream=True)\n",
    "                    print(response)\n",
    "                    # Check if the request was successful\n",
    "                    if response.status_code in [200, 206]:\n",
    "                        # Iterate over the content in chunks\n",
    "                        for chunk in response.iter_content(chunk_size=1024):\n",
    "                            if chunk:\n",
    "                                # Write the chunk to the file\n",
    "                                f.write(chunk)\n",
    "                                # Flush and sync the file to disk\n",
    "                                f.flush()\n",
    "                                os.fsync(f.fileno())\n",
    "                            progress_bar.update(len(chunk))\n",
    "                    else:\n",
    "                        print(f\"Failed to download file. Status code: {response.status_code}\")\n",
    "                except requests.exceptions.RequestException as e:\n",
    "                    print(\"Connection interrupted during download, resuming download from breaking point\")\n",
    "                    time.sleep(5)  # Sleep for a while before retrying the request\n",
    "                    continue  # Continue the loop to retry the request"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# View the checkpoints generated for the training job and for automl jobs, in addition view: best performing model's config and the results of all automl experiments\n",
    "\n",
    "if download_jobs:\n",
    "    if automl_enabled:\n",
    "        !python3 -m pip install pandas==1.5.1\n",
    "        import pandas as pd\n",
    "        model_downloaded_path = f\"{model_downloaded_path}/best_model\"\n",
    "        assert glob.glob(f\"{model_downloaded_path}/*.protobuf\") or glob.glob(f\"{model_downloaded_path}/*.yaml\")\n",
    "\n",
    "    assert os.path.exists(model_downloaded_path)\n",
    "    assert (glob.glob(model_downloaded_path + \"/**/*.tlt\", recursive=True) + glob.glob(model_downloaded_path + \"/**/*.hdf5\", recursive=True) + glob.glob(model_downloaded_path + \"/**/*.pth\", recursive=True))\n",
    "\n",
    "    if os.path.exists(model_downloaded_path):        \n",
    "        #List the binary model file\n",
    "        print(\"\\nCheckpoints for the training experiment\")\n",
    "        if os.path.exists(model_downloaded_path+\"/train/weights\") and len(os.listdir(model_downloaded_path+\"/train/weights\")) > 0:\n",
    "            print(f\"Folder: {model_downloaded_path}/train/weights\")\n",
    "            print(\"Files:\", os.listdir(model_downloaded_path+\"/train/weights\"))\n",
    "        elif os.path.exists(model_downloaded_path+\"/weights\") and len(os.listdir(model_downloaded_path+\"/weights\")) > 0:\n",
    "            print(f\"Folder: {model_downloaded_path}/weights\")\n",
    "            print(\"Files:\", os.listdir(model_downloaded_path+\"/weights\"))\n",
    "        else:\n",
    "            print(f\"Folder: {model_downloaded_path}\")\n",
    "            print(\"Files:\", os.listdir(model_downloaded_path))\n",
    "\n",
    "        if automl_enabled:\n",
    "            assert glob.glob(f\"{model_downloaded_path}/*.protobuf\") or glob.glob(f\"{model_downloaded_path}/*.yaml\")\n",
    "            experiment_artifacts = json.load(open(f\"{model_downloaded_path}/controller.json\",\"r\"))\n",
    "            data_frame = pd.DataFrame(experiment_artifacts)\n",
    "            # Print experiment id/number and the corresponding result\n",
    "            print(\"\\nResults of all experiments\")\n",
    "            with pd.option_context('display.max_rows', None, 'display.max_columns', None, 'display.max_colwidth', None):\n",
    "                print(data_frame[[\"id\",\"result\"]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Evaluate <a class=\"anchor\" id=\"head-12\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Get model handler parameters\n",
    "endpoint = f\"{base_url}/experiments/{experiment_id}\"\n",
    "response = requests.get(endpoint, headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "assert response.json()\n",
    "\n",
    "model_parameters = response.json()\n",
    "update_checkpoint_choosing = {}\n",
    "update_checkpoint_choosing[\"checkpoint_choose_method\"] = model_parameters[\"checkpoint_choose_method\"]\n",
    "update_checkpoint_choosing[\"checkpoint_epoch_number\"] = model_parameters[\"checkpoint_epoch_number\"]\n",
    "print(update_checkpoint_choosing)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Change the method by which checkpoint from the parent action is chosen, when parent action is a train/retrain action.\n",
    "# Example for evaluate action below, can be applied in the same way for other actions too\n",
    "update_checkpoint_choosing[\"checkpoint_choose_method\"] = \"latest_model\" # Choose between best_model/latest_model/from_epoch_number\n",
    "# If from_epoch_number is chosen then assign the epoch number to the dictionary key in the format 'from_epoch_number{train_job_id}'\n",
    "# update_checkpoint_choosing[\"checkpoint_epoch_number\"][\"from_epoch_number_28a2754e-50ef-43a8-9733-98913776dd90\"] = 3\n",
    "data = json.dumps(update_checkpoint_choosing)\n",
    "\n",
    "endpoint = f\"{base_url}/experiments/{experiment_id}\"\n",
    "\n",
    "response = requests.patch(endpoint, data=data, headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "print(response)\n",
    "print(json.dumps(response.json(), sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Get default spec schema\n",
    "endpoint = f\"{base_url}/experiments/{experiment_id}/specs/evaluate/schema\"\n",
    "\n",
    "response = requests.get(endpoint, headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "\n",
    "print(response)\n",
    "#print(response.json()) ## Uncomment for verbose schema\n",
    "assert \"default\" in response.json().keys()\n",
    "eval_specs = response.json()[\"default\"]\n",
    "print(json.dumps(eval_specs, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Apply changes\n",
    "if model_name == \"action_recognition\":\n",
    "    eval_specs[\"model\"][\"model_type\"] = model_type\n",
    "    eval_specs[\"model\"][\"input_type\"] = model_input_type\n",
    "elif model_name == \"fpenet\":\n",
    "    eval_specs[\"dataloader\"][\"dataset_info\"][\"root_path\"] = None\n",
    "    eval_specs[\"num_keypoints\"] = int(model_type)\n",
    "    eval_specs[\"dataloader\"][\"num_keypoints\"] = int(model_type)\n",
    "elif model_name == \"pose_classification\":\n",
    "    if model_type == \"nvidia\":\n",
    "        eval_specs[\"dataset\"][\"num_classes\"] = 6\n",
    "        eval_specs[\"model\"][\"graph_layout\"] = \"nvidia\"\n",
    "    elif model_type == \"kinetics\":\n",
    "        eval_specs[\"dataset\"][\"num_classes\"] = 5\n",
    "        eval_specs[\"model\"][\"graph_layout\"] = \"openpose\"\n",
    "elif model_name == \"re_identification\":\n",
    "    eval_specs[\"dataset\"][\"num_classes\"] = 100 #The number set in obtain_subset script\n",
    "elif model_name == \"visual_changenet\" and ds_format == 'visual_changenet_segment':\n",
    "    eval_specs[\"task\"] = 'segment'\n",
    "elif model_name == \"visual_changenet\" and ds_format == 'visual_changenet_classify':\n",
    "    eval_specs[\"task\"] = 'classify'\n",
    "    eval_specs[\"train\"][\"classify\"][\"loss\"] = \"contrastive\"\n",
    "elif model_name == \"centerpose\":\n",
    "    eval_specs[\"dataset\"][\"category\"] = testing_categories\n",
    "print(json.dumps(eval_specs, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Run action\n",
    "if model_name == \"visual_changenet\":\n",
    "    parent = job_map[\"train_\" + ds_format]\n",
    "else:\n",
    "    parent = job_map[\"train_\" + model_name]\n",
    "action = \"evaluate\"\n",
    "data = json.dumps({\"parent_job_id\":parent,\"action\":action,\"specs\":eval_specs})\n",
    "\n",
    "endpoint = f\"{base_url}/experiments/{experiment_id}/jobs\"\n",
    "\n",
    "response = requests.post(endpoint, data=data, headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "assert response.json()\n",
    "\n",
    "print(response)\n",
    "print(response.json())\n",
    "\n",
    "if model_name == \"visual_changenet\":\n",
    "    job_map[\"evaluate_\" + ds_format] = response.json()\n",
    "else:\n",
    "    job_map[\"evaluate_\" + model_name] = response.json()\n",
    "print(job_map)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Monitor job status by repeatedly running this cell\n",
    "if model_name == \"visual_changenet\":\n",
    "    job_id = job_map[\"evaluate_\" + ds_format]\n",
    "else:\n",
    "    job_id = job_map[\"evaluate_\" + model_name]\n",
    "endpoint = f\"{base_url}/experiments/{experiment_id}/jobs/{job_id}\"\n",
    "\n",
    "while True:\n",
    "    clear_output(wait=True)\n",
    "    response = requests.get(endpoint, headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "    print(response)\n",
    "    print(response.json())\n",
    "    assert \"status\" in response.json().keys() and response.json().get(\"status\") != \"Error\"\n",
    "    if response.json().get(\"status\") in [\"Done\",\"Error\", \"Canceled\"] or response.status_code not in (200,201):\n",
    "        break\n",
    "    time.sleep(15)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prune, Retrain and Evaluation <a class=\"anchor\" id=\"head-13\"></a>\n",
    "\n",
    "- We optimize the trained model by pruning and retraining in the following cells"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Prune <a class=\"anchor\" id=\"head-14\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Get default spec schema\n",
    "if model_name in (\"bpnet\", \"ocdnet\", \"ocrnet\", \"pointpillars\"):\n",
    "    endpoint = f\"{base_url}/experiments/{experiment_id}/specs/prune/schema\"\n",
    "\n",
    "    response = requests.get(endpoint, headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "\n",
    "    print(response)\n",
    "    #print(response.json()) ## Uncomment for verbose schema\n",
    "    assert \"default\" in response.json().keys()\n",
    "    prune_specs = response.json()[\"default\"]\n",
    "    print(json.dumps(prune_specs, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Apply changes\n",
    "# None for prune\n",
    "if model_name in (\"bpnet\", \"ocdnet\", \"ocrnet\", \"pointpillars\"):\n",
    "    print(json.dumps(prune_specs, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run actions\n",
    "if model_name in (\"bpnet\", \"ocdnet\", \"ocrnet\", \"pointpillars\"):\n",
    "    parent = job_map[\"train_\" + model_name]\n",
    "    action = \"prune\"\n",
    "    data = json.dumps({\"parent_job_id\":parent,\"action\":action,\"specs\":prune_specs})\n",
    "\n",
    "    endpoint = f\"{base_url}/experiments/{experiment_id}/jobs\"\n",
    "\n",
    "    response = requests.post(endpoint, data=data, headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "    assert response.json()\n",
    "\n",
    "    print(response)\n",
    "    print(response.json())\n",
    "\n",
    "    job_map[\"prune_\" + model_name] = response.json()\n",
    "    print(job_map)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Monitor job status by repeatedly running this cell (prune)\n",
    "if model_name in (\"bpnet\", \"ocdnet\", \"ocrnet\", \"pointpillars\"):\n",
    "    job_id = job_map[\"prune_\" + model_name]\n",
    "    endpoint = f\"{base_url}/experiments/{experiment_id}/jobs/{job_id}\"\n",
    "\n",
    "    while True:\n",
    "        clear_output(wait=True)\n",
    "        response = requests.get(endpoint, headers=headers)\n",
    "        assert response.status_code in (200, 201)\n",
    "        print(response)\n",
    "        print(response.json())\n",
    "        assert \"status\" in response.json().keys() and response.json().get(\"status\") != \"Error\"\n",
    "        if response.json().get(\"status\") in [\"Done\",\"Error\", \"Canceled\"] or response.status_code not in (200,201):\n",
    "            break\n",
    "        time.sleep(15)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Retrain <a class=\"anchor\" id=\"head-15\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Get default spec schema\n",
    "if model_name in (\"bpnet\", \"ocdnet\", \"ocrnet\", \"pointpillars\"):\n",
    "    endpoint = f\"{base_url}/experiments/{experiment_id}/specs/retrain/schema\"\n",
    "\n",
    "    response = requests.get(endpoint,headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "\n",
    "    print(response)\n",
    "    #print(response.json()) ## Uncomment for verbose schema\n",
    "    assert \"default\" in response.json().keys()\n",
    "    retrain_specs = response.json()[\"default\"]\n",
    "    print(json.dumps(retrain_specs, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Apply changes for any of the parameters listed in the previous cell as required\n",
    "if model_name in (\"bpnet\", \"ocdnet\", \"ocrnet\", \"pointpillars\"):\n",
    "    if model_name == \"bpnet\":\n",
    "        retrain_specs[\"num_epoch\"] = 20\n",
    "        retrain_specs[\"checkpoint_n_epoch\"] = 5\n",
    "        retrain_specs[\"validation_every_n_epoch\"] = 5\n",
    "        retrain_specs[\"finetuning_config\"][\"checkpoint_path\"] = None\n",
    "        retrain_specs[\"gpus\"] = 1\n",
    "    elif model_name == \"ocdnet\":\n",
    "        retrain_specs[\"train\"][\"num_epochs\"] = 30\n",
    "        retrain_specs[\"train\"][\"checkpoint_interval\"] = 5\n",
    "        retrain_specs[\"train\"][\"validation_interval\"] = 5\n",
    "        retrain_specs[\"train\"][\"gpu_id\"] = [0]\n",
    "        retrain_specs[\"num_gpus\"] = 1\n",
    "    elif model_name == \"ocrnet\":\n",
    "        retrain_specs[\"train\"][\"num_epochs\"] = 20\n",
    "        retrain_specs[\"train\"][\"checkpoint_interval\"] = 5\n",
    "        retrain_specs[\"train\"][\"validation_interval\"] = 5\n",
    "        retrain_specs[\"train\"][\"num_gpus\"] = 1\n",
    "    elif model_name == \"pointpillars\":\n",
    "        retrain_specs[\"train\"][\"num_epochs\"] = 80\n",
    "        retrain_specs[\"gpus\"] = 1\n",
    "    print(json.dumps(retrain_specs, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run actions\n",
    "if model_name in (\"bpnet\", \"ocdnet\", \"ocrnet\", \"pointpillars\"):\n",
    "    parent = job_map[\"prune_\" + model_name]\n",
    "    action = \"retrain\"\n",
    "    data = json.dumps({\"parent_job_id\":parent,\"action\":action,\"specs\":retrain_specs})\n",
    "\n",
    "    endpoint = f\"{base_url}/experiments/{experiment_id}/jobs\"\n",
    "\n",
    "    response = requests.post(endpoint, data=data, headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "    assert response.json()\n",
    "\n",
    "    print(response)\n",
    "    print(response.json())\n",
    "\n",
    "    job_map[\"retrain_\" + model_name] = response.json()\n",
    "    print(job_map)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Monitor job status by repeatedly running this cell (retrain)\n",
    "if model_name in (\"bpnet\", \"ocdnet\", \"ocrnet\", \"pointpillars\"):\n",
    "    job_id = job_map[\"retrain_\" + model_name]\n",
    "    endpoint = f\"{base_url}/experiments/{experiment_id}/jobs/{job_id}\"\n",
    "\n",
    "    while True:\n",
    "        clear_output(wait=True)\n",
    "        response = requests.get(endpoint, headers=headers)\n",
    "        assert response.status_code in (200, 201)\n",
    "        print(response)\n",
    "        print(response.json())\n",
    "        assert \"status\" in response.json().keys() and response.json().get(\"status\") != \"Error\"\n",
    "        if response.json().get(\"status\") in [\"Done\",\"Error\", \"Canceled\"] or response.status_code not in (200,201):\n",
    "            break\n",
    "        time.sleep(15)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Optional cancel job - for jobs that are pending/running (retrain)\n",
    "\n",
    "# if model_name == \"pointpillars\":\n",
    "#     job_id = job_map[\"retrain_\" + model_name]\n",
    "#     endpoint = f\"{base_url}/experiments/{experiment_id}/jobs/{job_id}:cancel\"\n",
    "\n",
    "#     response = requests.post(endpoint, headers=headers)\n",
    "#     assert response.status_code in (200, 201)\n",
    "\n",
    "#     print(response)\n",
    "#     print(response.json())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Optional delete job - for jobs that are error/done (retrain)\n",
    "\n",
    "# if model_name == \"pointpillars\":\n",
    "#     job_id = job_map[\"retrain_\" + model_name]\n",
    "#     endpoint = f\"{base_url}/experiments/{experiment_id}/jobs/{job_id}\"\n",
    "\n",
    "#     response = requests.delete(endpoint, headers=headers)\n",
    "#     assert response.status_code in (200, 201)\n",
    "\n",
    "#     print(response)\n",
    "#     print(response.json())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Evaluate after retrain <a class=\"anchor\" id=\"head-15\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get default spec schema\n",
    "if model_name in (\"bpnet\", \"ocdnet\", \"ocrnet\", \"pointpillars\"):\n",
    "    endpoint = f\"{base_url}/experiments/{experiment_id}/specs/evaluate/schema\"\n",
    "\n",
    "    response = requests.get(endpoint, headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "\n",
    "    print(response)\n",
    "    #print(response.json()) ## Uncomment for verbose schema\n",
    "    assert \"default\" in response.json().keys()\n",
    "    eval_retrain_specs = response.json()[\"default\"]\n",
    "    print(json.dumps(eval_retrain_specs, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Apply changes to specs if necessary\n",
    "if model_name in (\"bpnet\", \"ocdnet\", \"ocrnet\", \"pointpillars\"):\n",
    "    print(json.dumps(eval_retrain_specs, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run actions\n",
    "if model_name in (\"bpnet\", \"ocdnet\", \"ocrnet\", \"pointpillars\"):\n",
    "    parent = job_map[\"retrain_\" + model_name]\n",
    "    action = \"evaluate\"\n",
    "    data = json.dumps({\"parent_job_id\":parent,\"action\":action,\"specs\":eval_retrain_specs})\n",
    "\n",
    "    endpoint = f\"{base_url}/experiments/{experiment_id}/jobs\"\n",
    "\n",
    "    response = requests.post(endpoint, data=data, headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "    assert response.json()\n",
    "\n",
    "    print(response)\n",
    "    print(response.json())\n",
    "\n",
    "    job_map[\"eval_retrain_\" + model_name] = response.json()\n",
    "    print(job_map)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Monitor job status by repeatedly running this cell (evaluate)\n",
    "if model_name in (\"bpnet\", \"ocdnet\", \"ocrnet\", \"pointpillars\"):\n",
    "    job_id = job_map[\"eval_retrain_\" + model_name]\n",
    "    endpoint = f\"{base_url}/experiments/{experiment_id}/jobs/{job_id}\"\n",
    "\n",
    "    while True:\n",
    "        clear_output(wait=True)\n",
    "        response = requests.get(endpoint, headers=headers)\n",
    "        assert response.status_code in (200, 201)\n",
    "        print(response)\n",
    "        print(response.json())\n",
    "        assert \"status\" in response.json().keys() and response.json().get(\"status\") != \"Error\"\n",
    "        if response.json().get(\"status\") in [\"Done\",\"Error\", \"Canceled\"] or response.status_code not in (200,201):\n",
    "            break\n",
    "        time.sleep(15)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Export <a class=\"anchor\" id=\"head-17\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Get default spec schema\n",
    "endpoint = f\"{base_url}/experiments/{experiment_id}/specs/export/schema\"\n",
    "\n",
    "response = requests.get(endpoint, headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "\n",
    "print(response)\n",
    "# print(response.json()) ## Uncomment for verbose schema\n",
    "assert \"default\" in response.json().keys()\n",
    "export_specs = response.json()[\"default\"]\n",
    "print(json.dumps(export_specs, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Apply changes to the export_specs dictionary if necessary\n",
    "if model_name == \"action_recognition\":\n",
    "    export_specs[\"model\"][\"model_type\"] = model_type\n",
    "    export_specs[\"model\"][\"input_type\"] = model_input_type\n",
    "elif model_name == \"bpnet\":\n",
    "    export_specs[\"data_type\"] = \"int8\"\n",
    "    export_specs[\"max_batch_size\"] = 1\n",
    "elif model_name == \"lprnet\":\n",
    "    export_specs[\"data_type\"] = \"fp32\"\n",
    "elif model_name == \"pose_classification\":\n",
    "    if model_type == \"nvidia\":\n",
    "        export_specs[\"dataset\"][\"num_classes\"] = 6\n",
    "        export_specs[\"model\"][\"graph_layout\"] = \"nvidia\"\n",
    "    elif model_type == \"kinetics\":\n",
    "        export_specs[\"dataset\"][\"num_classes\"] = 5\n",
    "        export_specs[\"model\"][\"graph_layout\"] = \"openpose\"\n",
    "elif model_name == \"re_identification\":\n",
    "    export_specs[\"dataset\"][\"num_classes\"] = 100 #The number set in obtain_subset script\n",
    "elif model_name == \"visual_changenet\" and ds_format == 'visual_changenet_segment':\n",
    "    export_specs[\"export\"][\"input_height\"] = 256 \n",
    "    export_specs[\"export\"][\"input_width\"] = 256 \n",
    "    export_specs[\"task\"] = 'segment'\n",
    "elif model_name == \"visual_changenet\" and ds_format == 'visual_changenet_classify':\n",
    "    export_specs[\"export\"][\"input_height\"] = 512 \n",
    "    export_specs[\"export\"][\"input_width\"] = 128\n",
    "    export_specs[\"task\"] = 'classify'\n",
    "print(json.dumps(export_specs, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Run action\n",
    "if model_name == \"visual_changenet\":\n",
    "    parent = job_map[\"train_\" + ds_format]\n",
    "else:\n",
    "    parent = job_map[\"train_\" + model_name]\n",
    "action = \"export\"\n",
    "data = json.dumps({\"parent_job_id\":parent,\"action\":action,\"specs\":export_specs})\n",
    "\n",
    "endpoint = f\"{base_url}/experiments/{experiment_id}/jobs\"\n",
    "\n",
    "response = requests.post(endpoint, data=data, headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "assert response.json()\n",
    "\n",
    "print(response)\n",
    "print(response.json())\n",
    "\n",
    "if model_name == \"visual_changenet\":\n",
    "    job_map[\"export_\" + ds_format] = response.json()\n",
    "else:\n",
    "    job_map[\"export_\" + model_name] = response.json()\n",
    "print(job_map)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Monitor job status by repeatedly running this cell\n",
    "if model_name == \"visual_changenet\":\n",
    "    job_id = job_map[\"export_\" + ds_format]\n",
    "else:\n",
    "    job_id = job_map[\"export_\" + model_name]\n",
    "endpoint = f\"{base_url}/experiments/{experiment_id}/jobs/{job_id}\"\n",
    "\n",
    "while True:\n",
    "    clear_output(wait=True)\n",
    "    response = requests.get(endpoint, headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "    print(response)\n",
    "    print(response.json())\n",
    "    assert \"status\" in response.json().keys() and response.json().get(\"status\") != \"Error\"\n",
    "    if response.json().get(\"status\") in [\"Done\",\"Error\", \"Canceled\"] or response.status_code not in (200,201):\n",
    "        break\n",
    "    time.sleep(15)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### TRT Engine generation using TAO-Deploy <a class=\"anchor\" id=\"head-18\"></a>\n",
    "\n",
    "- Here, we use the exported model to convert to target platform"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Get default spec schema\n",
    "if model_name in (\"bpnet\",\"lprnet\", \"ocdnet\", \"ocrnet\", \"optical_inspection\", \"ml_recog\", \"visual_changenet\", \"centerpose\"):\n",
    "    if model_name == \"bpnet\":\n",
    "        engine_generation_action = \"trtexec\"\n",
    "    else:\n",
    "        engine_generation_action = \"gen_trt_engine\"\n",
    "    endpoint = f\"{base_url}/experiments/{experiment_id}/specs/{engine_generation_action}/schema\"\n",
    "\n",
    "    response = requests.get(endpoint, headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "\n",
    "    print(response)\n",
    "    #print(response.json()) ## Uncomment for verbose schema\n",
    "    assert \"default\" in response.json().keys()\n",
    "    tao_deploy_specs = response.json()[\"default\"]\n",
    "    print(json.dumps(tao_deploy_specs, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Apply changes\n",
    "if model_name in (\"bpnet\",\"lprnet\", \"ocdnet\", \"ocrnet\", \"optical_inspection\", \"ml_recog\", \"visual_changenet\", \"centerpose\"):\n",
    "    if model_name == \"lprnet\":\n",
    "        tao_deploy_specs[\"data_type\"] = \"fp32\"\n",
    "    elif model_name in (\"ml_recog\", \"ocdnet\"):\n",
    "        tao_deploy_specs[\"gen_trt_engine\"][\"tensorrt\"][\"data_type\"] = \"int8\"\n",
    "    elif model_name in (\"ocrnet\", \"optical_inspection\"):\n",
    "        tao_deploy_specs[\"gen_trt_engine\"][\"tensorrt\"][\"data_type\"] = \"fp16\"\n",
    "    elif model_name == \"visual_changenet\" and ds_format == 'visual_changenet_classify':\n",
    "        tao_deploy_specs[\"gen_trt_engine\"][\"input_height\"] = 512 \n",
    "        tao_deploy_specs[\"gen_trt_engine\"][\"input_width\"] = 128\n",
    "        tao_deploy_specs[\"task\"] = 'classify'\n",
    "    elif model_name == \"visual_changenet\" and ds_format == 'visual_changenet_segment':\n",
    "        tao_deploy_specs[\"gen_trt_engine\"][\"tensorrt\"][\"data_type\"] = \"fp16\"\n",
    "        tao_deploy_specs[\"gen_trt_engine\"][\"input_height\"] = 256\n",
    "        tao_deploy_specs[\"gen_trt_engine\"][\"input_width\"]= 256\n",
    "        tao_deploy_specs[\"task\"] = 'segment'\n",
    "    print(json.dumps(tao_deploy_specs, sort_keys=True, indent=4))        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Run action\n",
    "if model_name in (\"bpnet\", \"lprnet\", \"ocdnet\", \"ocrnet\", \"optical_inspection\", \"ml_recog\", \"visual_changenet\", \"centerpose\"):\n",
    "    if model_name == \"visual_changenet\":\n",
    "        parent = job_map[\"export_\" + ds_format]\n",
    "    else:\n",
    "        parent = job_map[\"export_\" + model_name]\n",
    "    action = engine_generation_action\n",
    "    data = json.dumps({\"parent_job_id\":parent,\"action\":action,\"specs\":tao_deploy_specs})\n",
    "\n",
    "    endpoint = f\"{base_url}/experiments/{experiment_id}/jobs\"\n",
    "\n",
    "    response = requests.post(endpoint, data=data, headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "    assert response.json()\n",
    "\n",
    "    print(response)\n",
    "    print(response.json())\n",
    "\n",
    "    if model_name == \"visual_changenet\":\n",
    "        job_map[\"gen_trt_engine_\" + ds_format] = response.json()\n",
    "    else:\n",
    "        job_map[\"gen_trt_engine_\" + model_name] = response.json()\n",
    "    print(job_map)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Monitor job status by repeatedly running this cell\n",
    "if model_name in (\"bpnet\", \"lprnet\", \"ocdnet\", \"ocrnet\", \"optical_inspection\", \"ml_recog\", \"visual_changenet\", \"centerpose\"):\n",
    "    if model_name == \"visual_changenet\":\n",
    "        job_id = job_map[\"gen_trt_engine_\" + ds_format]\n",
    "    else:\n",
    "        job_id = job_map[\"gen_trt_engine_\" + model_name]\n",
    "    endpoint = f\"{base_url}/experiments/{experiment_id}/jobs/{job_id}\"\n",
    "\n",
    "    while True:    \n",
    "        clear_output(wait=True)\n",
    "        response = requests.get(endpoint, headers=headers)\n",
    "        assert response.status_code in (200, 201)\n",
    "        print(response)\n",
    "        print(response.json())\n",
    "        assert \"status\" in response.json().keys() and response.json().get(\"status\") != \"Error\"\n",
    "        if response.json().get(\"status\") in [\"Done\",\"Error\", \"Canceled\"] or response.status_code not in (200,201):\n",
    "            break\n",
    "        time.sleep(15)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### TAO inference <a class=\"anchor\" id=\"head-19\"></a>\n",
    "\n",
    "- Run inference on a set of images using the .tlt model created at train step"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Get default spec schema\n",
    "endpoint = f\"{base_url}/experiments/{experiment_id}/specs/inference/schema\"\n",
    "\n",
    "response = requests.get(endpoint, headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "\n",
    "print(response)\n",
    "# print(response.json()) ## Uncomment for verbose schema\n",
    "assert \"default\" in response.json().keys()\n",
    "tao_inference_specs = response.json()[\"default\"]\n",
    "print(json.dumps(tao_inference_specs, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Apply changes to the tao_inference_specs dictionary if necessary\n",
    "if model_name == \"action_recognition\":\n",
    "    tao_inference_specs[\"model\"][\"model_type\"] = model_type\n",
    "    tao_inference_specs[\"model\"][\"input_type\"] = model_input_type\n",
    "elif model_name == \"fpenet\":\n",
    "    tao_inference_specs[\"num_keypoints\"] = int(model_type)\n",
    "    tao_inference_specs[\"dataloader\"][\"num_keypoints\"] = int(model_type)\n",
    "elif model_name == \"pose_classification\":\n",
    "    if model_type == \"nvidia\":\n",
    "        tao_inference_specs[\"dataset\"][\"num_classes\"] = 6\n",
    "        tao_inference_specs[\"model\"][\"graph_layout\"] = \"nvidia\"\n",
    "    elif model_type == \"kinetics\":\n",
    "        tao_inference_specs[\"dataset\"][\"num_classes\"] = 5\n",
    "        tao_inference_specs[\"model\"][\"graph_layout\"] = \"openpose\"\n",
    "elif model_name == \"re_identification\":\n",
    "    tao_inference_specs[\"dataset\"][\"num_classes\"] = 100 #The number set in obtain_subset script\n",
    "elif model_name == \"visual_changenet\" and ds_format == 'visual_changenet_classify':\n",
    "    tao_inference_specs[\"inference\"][\"batch_size\"] = tao_inference_specs[\"dataset\"][\"classify\"]['batch_size'] \n",
    "    tao_inference_specs[\"task\"] = 'classify'\n",
    "elif model_name == \"visual_changenet\" and ds_format == 'visual_changenet_segment':\n",
    "    tao_inference_specs[\"inference\"][\"batch_size\"] = tao_inference_specs[\"dataset\"][\"segment\"]['batch_size'] \n",
    "    tao_inference_specs[\"task\"] = 'segment'\n",
    "elif model_name == \"centerpose\":\n",
    "    tao_inference_specs[\"dataset\"][\"category\"] = testing_categories\n",
    "print(json.dumps(tao_inference_specs, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Run action\n",
    "if model_name == \"visual_changenet\":\n",
    "    parent = job_map[\"train_\" + ds_format]\n",
    "else:\n",
    "    parent = job_map[\"train_\" + model_name]\n",
    "action = \"inference\"\n",
    "data = json.dumps({\"parent_job_id\":parent,\"action\":action,\"specs\":tao_inference_specs})\n",
    "\n",
    "endpoint = f\"{base_url}/experiments/{experiment_id}/jobs\"\n",
    "\n",
    "response = requests.post(endpoint, data=data, headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "assert response.json()\n",
    "\n",
    "print(response)\n",
    "print(response.json())\n",
    "\n",
    "if model_name == \"visual_changenet\":\n",
    "    job_map[\"inference_tlt_\" + ds_format] = response.json()\n",
    "else:\n",
    "    job_map[\"inference_tlt_\" + model_name] = response.json()\n",
    "print(job_map)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Monitor job status by repeatedly running this cell\n",
    "if model_name == \"visual_changenet\":\n",
    "    job_id = job_map[\"inference_tlt_\" + ds_format]\n",
    "else:\n",
    "    job_id = job_map[\"inference_tlt_\" + model_name]\n",
    "endpoint = f\"{base_url}/experiments/{experiment_id}/jobs/{job_id}\"\n",
    "\n",
    "while True:\n",
    "    clear_output(wait=True)\n",
    "    response = requests.get(endpoint, headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "    print(response)\n",
    "    print(response.json())\n",
    "    assert \"status\" in response.json().keys() and response.json().get(\"status\") != \"Error\"\n",
    "    if response.json().get(\"status\") in [\"Done\",\"Error\", \"Canceled\"] or response.status_code not in (200,201):\n",
    "        break\n",
    "    time.sleep(15)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Download job contents once the above job shows \"Done\" status\n",
    "if download_jobs:\n",
    "    if model_name == \"visual_changenet\":\n",
    "        job_id = job_map[\"inference_tlt_\" + ds_format]\n",
    "    else:\n",
    "        job_id = job_map[\"inference_tlt_\" + model_name]\n",
    "    endpoint = f\"{base_url}/experiments/{experiment_id}/jobs/{job_id}\"\n",
    "    response = requests.get(endpoint, headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "    expected_file_size = response.json().get(\"job_tar_stats\", {}).get(\"file_size\")\n",
    "    print(\"expected_file_size: \", expected_file_size)\n",
    "\n",
    "    !python3 -m pip install tqdm\n",
    "    from tqdm import tqdm\n",
    "\n",
    "    endpoint = f'{base_url}/experiments/{experiment_id}/jobs/{job_id}:download'\n",
    "    temptar = f'{job_id}.tar.gz'\n",
    "\n",
    "    with tqdm(total=expected_file_size, unit='B', unit_scale=True) as progress_bar:\n",
    "        while True:\n",
    "            # Check if the file already exists\n",
    "            headers_download_job = dict(headers)\n",
    "            if os.path.exists(temptar):\n",
    "                # Get the current file size\n",
    "                file_size = os.path.getsize(temptar)\n",
    "                print(f\"File size of dowloaded content until now is {file_size}\")\n",
    "\n",
    "                # If the file size matches the expected size, break out of the loop\n",
    "                if file_size >= (expected_file_size-1):\n",
    "                    print(\"Download completed successfully.\")\n",
    "                    print(\"Untarring\")\n",
    "                    # Untar to destination\n",
    "                    tar_command = f'tar -xf {temptar} -C {workdir}/'\n",
    "                    os.system(tar_command)\n",
    "                    os.remove(temptar)\n",
    "                    print(f\"Results at {workdir}/{job_id}\")\n",
    "                    inference_out_path = f\"{workdir}/{job_id}\"\n",
    "                    break\n",
    "\n",
    "                # Set the headers to resume the download from where it left off\n",
    "                headers_download_job['Range'] = f'bytes={file_size}-'\n",
    "            # Open the file for writing in binary mode\n",
    "            with open(temptar, 'ab') as f:\n",
    "                try:\n",
    "                    response = requests.get(endpoint, headers=headers_download_job, stream=True)\n",
    "                    print(response)\n",
    "                    # Check if the request was successful\n",
    "                    if response.status_code in [200, 206]:\n",
    "                        # Iterate over the content in chunks\n",
    "                        for chunk in response.iter_content(chunk_size=1024):\n",
    "                            if chunk:\n",
    "                                # Write the chunk to the file\n",
    "                                f.write(chunk)\n",
    "                                # Flush and sync the file to disk\n",
    "                                f.flush()\n",
    "                                os.fsync(f.fileno())\n",
    "                            progress_bar.update(len(chunk))\n",
    "                    else:\n",
    "                        print(f\"Failed to download file. Status code: {response.status_code}\")\n",
    "                except requests.exceptions.RequestException as e:\n",
    "                    print(\"Connection interrupted during download, resuming download from breaking point\")\n",
    "                    time.sleep(5)  # Sleep for a while before retrying the request\n",
    "                    continue  # Continue the loop to retry the request"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Inference output must be here\n",
    "if download_jobs:\n",
    "    if model_name in (\"action_recognition\",\"lprnet\",\"ocrnet\"):\n",
    "        assert os.path.exists(f'{inference_out_path}/logs_from_toolkit.txt')\n",
    "        !cat {inference_out_path}/logs_from_toolkit.txt\n",
    "    elif model_name in (\"bpnet\",\"pointpillars\", \"centerpose\"):\n",
    "        if model_name == \"bpnet\":\n",
    "            assert glob.glob(f\"{inference_out_path}/images_annotated/*.png\")\n",
    "        elif model_name == \"pointpillars\":\n",
    "            assert glob.glob(f\"{inference_out_path}/infer/detected_boxes/*.png\")\n",
    "        elif model_name == \"centerpose\":\n",
    "            assert glob.glob(f\"{inference_out_path}/inference/*.png\")\n",
    "        !python3 -m pip install matplotlib\n",
    "        import matplotlib.pyplot as plt\n",
    "        import matplotlib.image as mpimg\n",
    "        if model_name == \"bpnet\":\n",
    "            sample_image = glob.glob(f\"{inference_out_path}/images_annotated/*.png\")[0]\n",
    "        elif model_name == \"pointpillars\":\n",
    "            sample_image = glob.glob(f\"{inference_out_path}/infer/detected_boxes/*.png\")[0]\n",
    "        elif model_name == \"centerpose\":\n",
    "            sample_image = glob.glob(f\"{inference_out_path}/inference/*.png\")[0]\n",
    "        def display_photo(path):\n",
    "            img = mpimg.imread(path)\n",
    "            plt.figure(figsize = (int(img.shape[0]/100)*2,int(img.shape[1]/100)*2))\n",
    "            plt.axis('off')\n",
    "            imgplot = plt.imshow(img, aspect='auto')\n",
    "            plt.show()\n",
    "        display_photo(sample_image)\n",
    "    elif model_name == \"fpenet\":\n",
    "        assert os.path.exists(f'{inference_out_path}/result.txt')\n",
    "        !cat {inference_out_path}/result.txt\n",
    "    elif model_name == \"ml_recog\":\n",
    "        assert os.path.exists(f'{inference_out_path}/inference/result.csv')\n",
    "        !cat {inference_out_path}/inference/result.csv\n",
    "    elif model_name == \"optical_inspection\":\n",
    "        assert os.path.exists(f'{inference_out_path}/inference/inference.csv')\n",
    "        !cat {inference_out_path}/inference/inference.csv\n",
    "    elif model_name == \"pose_classification\":\n",
    "        assert os.path.exists(f'{inference_out_path}/results.txt')\n",
    "        !cat {inference_out_path}/results.txt\n",
    "    elif model_name == \"re_identification\":\n",
    "        assert os.path.exists(f'{inference_out_path}/inference.json')\n",
    "        !cat {inference_out_path}/inference.json\n",
    "    elif model_name == \"visual_changenet\":\n",
    "        if ds_format == 'visual_changenet_classify':\n",
    "            assert os.path.exists(f'{inference_out_path}/inference/inference.csv')\n",
    "            !cat {inference_out_path}/inference/inference.csv\n",
    "        elif ds_format == 'visual_changenet_segment':\n",
    "            assert os.path.exists(f'{inference_out_path}/inference/status.json')\n",
    "            !cat {inference_out_path}/inference/status.json"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### TRT inference <a class=\"anchor\" id=\"head-20\"></a>\n",
    "\n",
    "- no need to change the specs since we already uploaded it at the tlt inference step"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Get default spec schema\n",
    "if model_name in (\"bpnet\", \"lprnet\", \"ocdnet\", \"ocrnet\", \"ml_recog\", \"optical_inspection\", \"visual_changenet\", \"centerpose\"):\n",
    "    endpoint = f\"{base_url}/experiments/{experiment_id}/specs/inference/schema\"\n",
    "    response = requests.get(endpoint, headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "\n",
    "    print(response)\n",
    "    # print(response.json()) ## Uncomment for verbose schema\n",
    "    assert \"default\" in response.json().keys()\n",
    "    trt_inference_specs = response.json()[\"default\"]\n",
    "    print(json.dumps(trt_inference_specs, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Apply changes to the specs dictionary if necessary\n",
    "if model_name in (\"bpnet\", \"lprnet\", \"ocdnet\", \"ocrnet\", \"ml_recog\", \"optical_inspection\", \"visual_changenet\", \"centerpose\"):\n",
    "    if model_name == \"visual_changenet\" and ds_format == 'visual_changenet_classify':\n",
    "        trt_inference_specs[\"inference\"][\"batch_size\"] = trt_inference_specs[\"dataset\"][\"classify\"]['batch_size']\n",
    "        trt_inference_specs[\"task\"] = 'classify'\n",
    "    elif model_name == \"visual_changenet\" and ds_format == 'visual_changenet_segment':\n",
    "        trt_inference_specs[\"inference\"][\"batch_size\"] = trt_inference_specs[\"dataset\"][\"segment\"]['batch_size']\n",
    "        trt_inference_specs[\"task\"] = 'segment'\n",
    "    print(json.dumps(trt_inference_specs, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Run action\n",
    "if model_name in (\"bpnet\", \"lprnet\", \"ocdnet\", \"ocrnet\", \"ml_recog\", \"optical_inspection\", \"visual_changenet\", \"centerpose\"):\n",
    "    if model_name == \"visual_changenet\":\n",
    "        parent = job_map[\"gen_trt_engine_\" + ds_format]\n",
    "    else:\n",
    "        parent = job_map[\"gen_trt_engine_\" + model_name]\n",
    "    action = \"inference\"\n",
    "    data = json.dumps({\"parent_job_id\":parent,\"action\":action,\"specs\":trt_inference_specs})\n",
    "\n",
    "    endpoint = f\"{base_url}/experiments/{experiment_id}/jobs\"\n",
    "\n",
    "    response = requests.post(endpoint, data=data, headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "    assert response.json()\n",
    "\n",
    "    print(response)\n",
    "    print(response.json())\n",
    "\n",
    "    if model_name == \"visual_changenet\":\n",
    "        job_map[\"inference_trt_\" + ds_format] = response.json()\n",
    "    else:\n",
    "        job_map[\"inference_trt_\" + model_name] = response.json()\n",
    "    print(job_map)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Monitor job status by repeatedly running this cell\n",
    "if model_name in (\"bpnet\", \"lprnet\", \"ocdnet\", \"ocrnet\", \"ml_recog\", \"optical_inspection\", \"visual_changenet\", \"centerpose\"):\n",
    "    if model_name == \"visual_changenet\":\n",
    "        job_id = job_map[\"inference_trt_\" + ds_format]\n",
    "    else:\n",
    "        job_id = job_map[\"inference_trt_\" + model_name]\n",
    "    endpoint = f\"{base_url}/experiments/{experiment_id}/jobs/{job_id}\"\n",
    "\n",
    "    while True:    \n",
    "        clear_output(wait=True)\n",
    "        response = requests.get(endpoint, headers=headers)\n",
    "        assert response.status_code in (200, 201)\n",
    "        print(response)\n",
    "        print(response.json())\n",
    "        assert \"status\" in response.json().keys() and response.json().get(\"status\") != \"Error\"\n",
    "        if response.json().get(\"status\") in [\"Done\",\"Error\", \"Canceled\"] or response.status_code not in (200,201):\n",
    "            break\n",
    "        time.sleep(15)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Download job contents once the above job shows \"Done\" status\n",
    "if download_jobs:\n",
    "    if model_name in (\"bpnet\", \"lprnet\", \"ocdnet\", \"ocrnet\", \"ml_recog\", \"optical_inspection\", \"visual_changenet\", \"centerpose\"):\n",
    "        if model_name == \"visual_changenet\":\n",
    "            job_id = job_map[\"inference_trt_\" + ds_format]\n",
    "        else:\n",
    "            job_id = job_map[\"inference_trt_\" + model_name]\n",
    "        endpoint = f\"{base_url}/experiments/{experiment_id}/jobs/{job_id}\"\n",
    "        response = requests.get(endpoint, headers=headers)\n",
    "        assert response.status_code in (200, 201)\n",
    "        expected_file_size = response.json().get(\"job_tar_stats\", {}).get(\"file_size\")\n",
    "        print(\"expected_file_size: \", expected_file_size)\n",
    "\n",
    "        !python3 -m pip install tqdm\n",
    "        from tqdm import tqdm\n",
    "\n",
    "        endpoint = f'{base_url}/experiments/{experiment_id}/jobs/{job_id}:download'\n",
    "        temptar = f'{job_id}.tar.gz'\n",
    "\n",
    "        with tqdm(total=expected_file_size, unit='B', unit_scale=True) as progress_bar:\n",
    "            while True:\n",
    "                # Check if the file already exists\n",
    "                headers_download_job = dict(headers)\n",
    "                if os.path.exists(temptar):\n",
    "                    # Get the current file size\n",
    "                    file_size = os.path.getsize(temptar)\n",
    "                    print(f\"File size of dowloaded content until now is {file_size}\")\n",
    "\n",
    "                    # If the file size matches the expected size, break out of the loop\n",
    "                    if file_size >= (expected_file_size-1):\n",
    "                        print(\"Download completed successfully.\")\n",
    "                        print(\"Untarring\")\n",
    "                        # Untar to destination\n",
    "                        tar_command = f'tar -xf {temptar} -C {workdir}/'\n",
    "                        os.system(tar_command)\n",
    "                        os.remove(temptar)\n",
    "                        print(f\"Results at {workdir}/{job_id}\")\n",
    "                        inference_out_path = f\"{workdir}/{job_id}\"\n",
    "                        break\n",
    "\n",
    "                    # Set the headers to resume the download from where it left off\n",
    "                    headers_download_job['Range'] = f'bytes={file_size}-'\n",
    "                # Open the file for writing in binary mode\n",
    "                with open(temptar, 'ab') as f:\n",
    "                    try:\n",
    "                        response = requests.get(endpoint, headers=headers_download_job, stream=True)\n",
    "                        print(response)\n",
    "                        # Check if the request was successful\n",
    "                        if response.status_code in [200, 206]:\n",
    "                            # Iterate over the content in chunks\n",
    "                            for chunk in response.iter_content(chunk_size=1024):\n",
    "                                if chunk:\n",
    "                                    # Write the chunk to the file\n",
    "                                    f.write(chunk)\n",
    "                                    # Flush and sync the file to disk\n",
    "                                    f.flush()\n",
    "                                    os.fsync(f.fileno())\n",
    "                                progress_bar.update(len(chunk))\n",
    "                        else:\n",
    "                            print(f\"Failed to download file. Status code: {response.status_code}\")\n",
    "                    except requests.exceptions.RequestException as e:\n",
    "                        print(\"Connection interrupted during download, resuming download from breaking point\")\n",
    "                        time.sleep(5)  # Sleep for a while before retrying the request\n",
    "                        continue  # Continue the loop to retry the request"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Inference output must be here\n",
    "if download_jobs:\n",
    "    if model_name in (\"bpnet\",\"lprnet\", \"ocdnet\", \"ocrnet\", \"ml_recog\", \"optical_inspection\", \"visual_changenet\", \"centerpose\"):\n",
    "        !ls {inference_out_path}/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Delete model <a class=\"anchor\" id=\"head-21\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "endpoint = f\"{base_url}/experiments/{experiment_id}\"\n",
    "\n",
    "response = requests.delete(endpoint,headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "\n",
    "print(response)\n",
    "print(response.json())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Delete dataset <a class=\"anchor\" id=\"head-21\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Delete train dataset <a class=\"anchor\" id=\"head-21\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "endpoint = f\"{base_url}/datasets/{train_dataset_id}\"\n",
    "\n",
    "response = requests.delete(endpoint,headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "\n",
    "print(response)\n",
    "print(response.json())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Delete val dataset <a class=\"anchor\" id=\"head-21\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if model_name in (\"lprnet\", \"ocdnet\", \"ocrnet\", \"optical_inspection\") or ds_format == 'visual_changenet_classify':\n",
    "    endpoint = f\"{base_url}/datasets/{eval_dataset_id}\"\n",
    "\n",
    "    response = requests.delete(endpoint,headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "\n",
    "    print(response)\n",
    "    print(response.json())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Delete test dataset <a class=\"anchor\" id=\"head-21\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if model_name in (\"lprnet\", \"optical_inspection\") or ds_format == 'visual_changenet_classify':\n",
    "    endpoint = f\"{base_url}/datasets/{test_dataset_id}\"\n",
    "\n",
    "    response = requests.delete(endpoint,headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "\n",
    "    print(response)\n",
    "    print(response.json())"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
